@baimingtao/lbs-agent 1.0.0

package/extensions/skills/optional/mlops/pytorch-fsdp/references/other.md

@@ -0,0 +1,3297 @@
+# Pytorch-Fsdp - 其他
+
+**页数：** 15
+
+---
+
+## 分布式数据并行#
+
+**URL:** https://pytorch.org/docs/stable/notes/ddp.html
+
+**目录：**
+- 分布式数据并行#
+- 示例#
+- 内部设计#
+- 实现#
+  - ProcessGroup#
+  - DistributedDataParallel#
+  - TorchDynamo DDPOptimizer#
+
+创建时间：2020年1月15日 | 最后更新时间：2024年1月25日
+
+torch.nn.parallel.DistributedDataParallel 的实现随着时间不断演进。本设计说明基于 v1.4 版本的状态编写。
+
+torch.nn.parallel.DistributedDataParallel (DDP) 透明地执行分布式数据并行训练。本页介绍了它的工作原理并揭示了实现细节。
+
+让我们从一个简单的 torch.nn.parallel.DistributedDataParallel 示例开始。此示例使用 torch.nn.Linear 作为本地模型，用 DDP 包装它，然后在 DDP 模型上运行一次前向传递、一次反向传递和一个优化器步骤。之后，本地模型上的参数将被更新，并且不同进程上的所有模型应该完全相同。
+
+DDP 可以与 TorchDynamo 配合使用。与 TorchDynamo 一起使用时，请在编译模型之前应用 DDP 模型包装器，以便 torchdynamo 可以基于 DDP 存储桶（bucket）大小应用 DDPOptimizer（计算图断点优化）。（有关更多信息，请参见 TorchDynamo DDPOptimizer。）
+
+本节通过深入探讨一次迭代中每个步骤的细节，揭示了 torch.nn.parallel.DistributedDataParallel 背后的工作原理。
+
+前提条件：DDP 依赖 c10d ProcessGroup 进行通信。因此，应用程序必须在构建 DDP 之前创建 ProcessGroup 实例。
+
+构建：DDP 构造函数接收对本地模块的引用，并将 state_dict() 从 rank 为 0 的进程广播到组中的所有其他进程，以确保所有模型副本从完全相同的状态开始。然后，每个 DDP 进程创建一个本地 Reducer，它稍后将在反向传递期间负责梯度同步。为了提高通信效率，Reducer 将参数梯度组织到存储桶中，并每次减少（reduce）一个存储桶。可以通过在 DDP 构造函数中设置 bucket_cap_mb 参数来配置存储桶大小。参数梯度到存储桶的映射是在构建时根据存储桶大小限制和参数大小确定的。模型参数以给定模型的 Model.parameters() 的（大致）逆序分配到存储桶中。使用逆序的原因是 DDP 预期在反向传递期间梯度大约按照该顺序准备就绪。下图显示了一个示例。请注意，grad0 和 grad1 在 bucket1 中，而另外两个梯度在 bucket0 中。当然，这个假设可能并不总是成立的，当这种情况发生时，可能会降低 DDP 反向传递的速度，因为 Reducer 无法在最早可能的时间启动通信。除了分桶之外，Reducer 还会在构建期间注册自动求导钩子，每个参数一个钩子。这些钩子将在反向传递期间梯度准备就绪时被触发。
+
+前向传递：DDP 接收输入并将其传递给本地模型，然后如果 find_unused_parameters 设置为 True，则会分析本地模型的输出。此模式允许在模型的子图上运行反向传递，DDP 通过从模型输出遍历自动求导图并将所有未使用的参数标记为已准备好进行归约（reduction）来找出哪些参数参与了反向传递。在反向传递期间，Reducer 只会等待未准备好的参数，但它仍会归约所有存储桶。将参数梯度标记为就绪目前并不能帮助 DDP 跳过存储桶，但它可以防止 DDP 在反向传递期间永远等待缺失的梯度。请注意，遍历自动求导图会引入额外的开销，因此应用程序应仅在必要时将 find_unused_parameters 设置为 True。
+
+反向传递：backward() 函数直接在损失 Tensor 上调用，这超出了 DDP 的控制范围，DDP 使用在构建时注册的自动求导钩子来触发梯度同步。当一个梯度准备就绪时，在该梯度累加器上相应的 DDP 钩子将触发，然后 DDP 会将该参数梯度标记为已准备好进行归约。当一个存储桶中的所有梯度都准备就绪时，Reducer 会启动该存储桶上的异步 allreduce 操作，以计算所有进程间梯度的平均值。当所有存储桶都准备就绪时，Reducer 将阻塞等待所有 allreduce 操作完成。完成后，平均梯度将被写入所有参数的 param.grad 字段。因此，在反向传递之后，不同 DDP 进程中相同对应参数上的 grad 字段应该是相同的。
+
+优化器步骤：从优化器的角度来看，它正在优化一个本地模型。所有 DDP 进程上的模型副本可以保持同步，因为它们都从相同的状态开始，并且在每次迭代中都具有相同的平均梯度。
+
+DDP 要求所有进程上的 Reducer 实例以完全相同的顺序调用 allreduce，这是通过始终按照存储桶索引顺序而不是实际的存储桶就绪顺序运行 allreduce 来实现的。跨进程的 allreduce 顺序不匹配可能导致错误结果或 DDP 反向传递挂起。
+
+以下是指向 DDP 实现组件的指针。堆叠图显示了代码的结构。
+
+ProcessGroup.hpp：包含所有进程组实现的抽象 API。c10d 库提供了 3 种开箱即用的实现，即 ProcessGroupGloo、ProcessGroupNCCL 和 ProcessGroupMPI。DistributedDataParallel 在初始化期间使用 ProcessGroup::broadcast() 将模型状态从 rank 为 0 的进程发送给其他进程，并使用 ProcessGroup::allreduce() 对梯度进行求和。
+
+Store.hpp：协助进程组实例互相发现的会合服务。
+
+distributed.py：是 DDP 的 Python 入口点。它实现了初始化步骤和 nn.parallel.DistributedDataParallel 模块的 forward 函数，这些函数调用了 C++ 库。当一个 DDP 进程在多个设备上工作时，它的 _sync_param 函数会执行进程内参数同步，并且它还会将模型缓冲区从 rank 为 0 的进程广播到所有其他进程。进程间的参数同步发生在 Reducer.cpp 中。
+
+comm.h：实现了合并广播辅助函数，该函数在初始化期间被调用来广播模型状态，并在前向传递之前同步模型缓冲区。
+
+reducer.h：提供了反向传递中梯度同步的核心实现。它有三个入口点函数：
+
+Reducer：在 distributed.py 中调用的构造函数，它将 Reducer::autograd_hook() 注册到梯度累加器。
+
+autograd_hook() 函数将在梯度准备就绪时由自动求导引擎调用。
+
+prepare_for_backward() 在 distributed.py 中的 DDP 前向传递结束时被调用。如果在 DDP 构造函数中 find_unused_parameters 设置为 True，它会遍历自动求导图以查找未使用的参数。
+
+DDP 的性能优势源于在反向传递期间将 allreduce 集合通信与计算重叠起来。当与 TorchDynamo 一起使用以编译整个前向和整个反向计算图时，AotAutograd 阻止了这种重叠，因为 allreduce 操作是在整个优化后的反向计算完成_之后_由自动求导钩子启动的。
+
+TorchDynamo 的 DDPOptimizer 通过在反向传递期间 DDP 的 allreduce 存储桶的逻辑边界处断开前向计算图来提供帮助。注意：目标是在反向传递期间断开计算图，最简单的实现是断开前向计算图，然后对每个部分调用 AotAutograd 和编译。这允许 DDP 的 allreduce 钩子在反向传递的各个部分之间触发，并调度通信以与计算重叠。
+
+有关更深入的解释和实验结果，请参阅此博客文章，或阅读 torch/_dynamo/optimizations/distributed.py 中的文档和代码
+
+要调试 DDPOptimizer，请设置 TORCH_LOGS='ddp_graphs' 以获取完整的计算图转储。对于不包含计算图的日志，请将 'dynamo'、'distributed' 或 'dist_ddp' 中的任何一个添加到 TORCH_LOGS 中（以获取有关存储桶边界的基本信息）。要禁用 DDPOptimizer，请设置 torch._dynamo.config.optimize_ddp=False。在没有 DDPOptimizer 的情况下，DDP 和 TorchDynamo 仍然应该能够正常工作，但性能会有所下降。
+
+---
+
+## PyTorch 文档#
+
+**URL:** https://pytorch.org/docs/stable/
+
+**目录：**
+- PyTorch 文档#
+- 索引和表格#
+
+PyTorch 是一个针对使用 GPU 和 CPU 的深度学习进行了优化的张量库。
+
+本档中描述的特性按发布状态进行分类：
+
+稳定版 (API-Stable)：这些功能将得到长期维护，并且通常在性能上没有重大限制或在文档上没有缺失。我们也期望保持向后兼容性（尽管可能会发生破坏性更改，并会提前一个版本发出通知）。
+
+不稳定版 (API-Unstable)：包含所有正在积极开发中的功能，这些 API 可能会根据用户反馈、必要的性能改进或由于跨算子的覆盖范围尚未完成而发生变化。这些功能的 API 和性能特征可能会发生变化。
+
+---
+
+## 通用 Join 上下文管理器#
+
+**URL:** https://pytorch.org/docs/stable/distributed.algorithms.join.html
+
+**目录：**
+- 通用 Join 上下文管理器#
+
+创建时间：2025年6月6日 | 最后更新时间：2025年6月6日
+
+通用 Join 上下文管理器有助于针对不均匀输入进行分布式训练。本页概述了相关类的 API：Join、Joinable 和 JoinHook。有关教程，请参阅使用 Join 上下文管理器进行不均匀输入的分布式训练。
+
+此类定义了通用的 join 上下文管理器，允许在某个进程加入（join）后调用自定义钩子。
+
+这些钩子应该掩盖（shadow）未加入进程的集合通信，以防止挂起和报错，并确保算法的正确性。有关钩子定义的详细信息，请参阅 JoinHook。
+
+上下文管理器要求每个参与的 Joinable 在其自身的每次迭代的集合通信之前调用 notify_join_context() 方法，以确保正确性。
+
+上下文管理器要求 JoinHook 对象中的所有 process_group 属性必须相同。如果有多个 JoinHook 对象，则使用第一个对象的设备。进程组和设备信息用于检查是否存在未加入的进程，并在启用 throw_on_early_termination 时通知进程抛出异常，这两者都使用 all-reduce 操作。
+
+joinables (List[Joinable]) – 参与的 Joinable 列表；它们的钩子将按给定顺序进行迭代。
+
+enable (bool) – 启用不均匀输入检测的标志；设置为 False 会禁用上下文管理器的功能，只能在用户确信输入不会不均匀时设置（默认：True）。
+
+throw_on_early_termination (bool) – 控制在检测到不均匀输入时是否抛出异常的标志（默认：False）。
+
+通知 join 上下文管理器，调用进程尚未加入。
+
+然后，如果 throw_on_early_termination=True，则检查是否检测到了不均匀输入（即是否有进程已经加入），如果是，则抛出异常。
+
+此方法应在其每次迭代的集合通信之前从 Joinable 对象中调用。例如，在 DistributedDataParallel 中，应在前向传递开始时调用此方法。
+
+只有传入上下文管理器的第一个 Joinable 对象会在此方法中执行集合通信，而对于其他对象，此方法是空操作。
+
+joinable (Joinable) – 调用此方法的 Joinable 对象。
+
+如果 joinable 是传入上下文管理器的第一个对象，则返回一个用于通知上下文管理器该进程尚未加入的 all-reduce 异步工作句柄；否则返回 None。
+
+这定义了可加入类的抽象基类。
+
+一个可加入类（继承自 Joinable）应该实现返回 JoinHook 实例的 join_hook()，此外还要分别返回设备和进程组信息的 join_device() 和 join_process_group()。
+
+返回执行 join 上下文管理器所需集合通信的设备。
+
+返回给定 Joinable 的 JoinHook 实例。
+
+kwargs (dict) – 一个包含在运行时修改 join 钩子行为的任何关键字参数的字典；所有共享同一个 join 上下文管理器的 Joinable 实例都会接收到相同的 kwargs 值。
+
+返回 join 上下文管理器本身所需集合通信的进程组。
+
+这定义了一个 join 钩子，它在 join 上下文管理器中提供了两个入口点。
+
+入口点：一个主钩子，只要存在未加入的进程就会被重复调用；一个后置钩子，在所有进程都加入后调用一次。
+
+要为通用 join 上下文管理器实现 join 钩子，请定义一个继承自 JoinHook 的类，并根据需要覆盖 main_hook() 和 post_hook()。
+
+只要存在未加入的进程，就调用此钩子以掩盖一次训练迭代中的集合通信。
+
+训练迭代，即在一次前向传递、反向传递和优化器步骤中。
+
+在所有进程都加入后调用此钩子。
+
+它被传递了一个额外的布尔参数 is_last_joiner，指示该 rank 是否是最后加入的之一。
+
+is_last_joiner (bool) – 如果该 rank 是最后加入的之一，则为 True；否则为 False。
+
+---
+
+## 实验性面向对象的分布式 API#
+
+**URL:** https://pytorch.org/docs/stable/distributed._dist2.html
+
+**目录：**
+- 实验性面向对象的分布式 API#
+
+创建时间：2025年7月9日 | 最后更新时间：2025年7月30日
+
+这是一个用于 PyTorch 分布式的全新实验性 API。目前正在积极开发中，可能会发生变更或被完全删除。
+
+这旨在作为更加灵活和面向对象的分布式 API 的试验场。
+
+基类：pybind11_object
+
+ProcessGroup 是一种通信原语，允许在一组进程之间进行集合操作。
+
+这是一个为所有 ProcessGroup 提供接口的基类。它不适合直接使用，而应由子类扩展。
+
+基类：pybind11_object
+
+用于进程组的后端类型。
+
+如果后端支持，中止所有操作和连接
+
+allgather(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[collections.abc.Sequence[torch.Tensor]], input_tensors: collections.abc.Sequence[torch.Tensor], opts: torch._C._distributed_c10d.AllgatherOptions = <torch._C._distributed_c10d.AllgatherOptions object at 0x7f0162b6b9b0>) -> c10d::Work
+
+跨进程组从所有进程中 allgather（全收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_gather()。
+
+allgather(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[torch.Tensor], input_tensor: torch.Tensor, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组从所有进程中 allgather（全收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_gather()。
+
+跨进程组从所有进程中 allgather（全收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_gather()。
+
+跨进程组从所有进程中 allgather（全收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_gather()。
+
+allreduce(self: torch._C._distributed_c10d.ProcessGroup, tensors: collections.abc.Sequence[torch.Tensor], opts: torch._C._distributed_c10d.AllreduceOptions = <torch._C._distributed_c10d.AllreduceOptions object at 0x7f0162745db0>) -> c10d::Work
+
+跨进程组中的所有进程对提供的张量进行 allreduce（全归约）。
+
+有关详细信息，请参见 torch.distributed.all_reduce()。
+
+allreduce(self: torch._C._distributed_c10d.ProcessGroup, tensors: collections.abc.Sequence[torch.Tensor], op: torch._C._distributed_c10d.ReduceOp = <RedOpType.SUM: 0>, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组中的所有进程对提供的张量进行 allreduce（全归约）。
+
+有关详细信息，请参见 torch.distributed.all_reduce()。
+
+allreduce(self: torch._C._distributed_c10d.ProcessGroup, tensor: torch.Tensor, op: torch._C._distributed_c10d.ReduceOp = <RedOpType.SUM: 0>, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组中的所有进程对提供的张量进行 allreduce（全归约）。
+
+有关详细信息，请参见 torch.distributed.all_reduce()。
+
+跨进程组中的所有进程对提供的张量进行 allreduce（全归约）。
+
+有关详细信息，请参见 torch.distributed.all_reduce()。
+
+跨进程组从所有进程中进行 input tensor（输入张量）的 alltoall（全交换）。
+
+有关详细信息，请参见 torch.distributed.all_to_all()。
+
+alltoall_base(self: torch._C._distributed_c10d.ProcessGroup, output: torch.Tensor, input: torch.Tensor, output_split_sizes: collections.abc.Sequence[typing.SupportsInt], input_split_sizes: collections.abc.Sequence[typing.SupportsInt], opts: torch._C._distributed_c10d.AllToAllOptions = <torch._C._distributed_c10d.AllToAllOptions object at 0x7f0162b79d30>) -> c10d::Work
+
+跨进程组从所有进程中 alltoall（全交换）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_to_all()。
+
+alltoall_base(self: torch._C._distributed_c10d.ProcessGroup, output: torch.Tensor, input: torch.Tensor, output_split_sizes: collections.abc.Sequence[typing.SupportsInt], input_split_sizes: collections.abc.Sequence[typing.SupportsInt], timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组从所有进程中 alltoall（全交换）输入张量。
+
+有关详细信息，请参见 torch.distributed.all_to_all()。
+
+barrier(self: torch._C._distributed_c10d.ProcessGroup, opts: torch._C._distributed_c10d.BarrierOptions = <torch._C._distributed_c10d.BarrierOptions object at 0x7f0162745ab0>) -> c10d::Work
+
+然后所有进程一起离开调用。
+
+有关详细信息，请参见 torch.distributed.barrier()。
+
+barrier(self: torch._C._distributed_c10d.ProcessGroup, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+然后所有进程一起离开调用。
+
+有关详细信息，请参见 torch.distributed.barrier()。
+
+broadcast(self: torch._C._distributed_c10d.ProcessGroup, tensors: collections.abc.Sequence[torch.Tensor], opts: torch._C._distributed_c10d.BroadcastOptions = <torch._C._distributed_c10d.BroadcastOptions object at 0x7f0162b7afb0>) -> c10d::Work
+
+将张量广播到进程组中的所有进程。
+
+有关详细信息，请参见 torch.distributed.broadcast()。
+
+broadcast(self: torch._C._distributed_c10d.ProcessGroup, tensor: torch.Tensor, root: typing.SupportsInt, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+将张量广播到进程组中的所有进程。
+
+有关详细信息，请参见 torch.distributed.broadcast()。
+
+gather(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[collections.abc.Sequence[torch.Tensor]], input_tensors: collections.abc.Sequence[torch.Tensor], opts: torch._C._distributed_c10d.GatherOptions = <torch._C._distributed_c10d.GatherOptions object at 0x7f0162c301f0>) -> c10d::Work
+
+跨进程组从所有进程中 gather（收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.gather()。
+
+gather(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[torch.Tensor], input_tensor: torch.Tensor, root: typing.SupportsInt, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组从所有进程中 gather（收集）输入张量。
+
+有关详细信息，请参见 torch.distributed.gather()。
+
+获取此进程组的存储。
+
+获取此进程组的描述
+
+（获取此进程组的名称。它在集群中是唯一的）
+
+然后所有进程一起离开调用。
+
+有关详细信息，请参见 torch.distributed.monitored_barrier()。
+
+获取此进程组的名称。
+
+获取此进程组的 rank。
+
+从指定的 rank 接收张量。
+
+有关详细信息，请参见 torch.distributed.recv()。
+
+从任何来源接收张量。
+
+有关详细信息，请参见 torch.distributed.recv()。
+
+reduce(self: torch._C._distributed_c10d.ProcessGroup, tensors: collections.abc.Sequence[torch.Tensor], opts: torch._C._distributed_c10d.ReduceOptions = <torch._C._distributed_c10d.ReduceOptions object at 0x7f0162bce3f0>) -> c10d::Work
+
+跨进程组中的所有进程归约提供的张量。
+
+有关详细信息，请参见 torch.distributed.reduce()。
+
+reduce(self: torch._C._distributed_c10d.ProcessGroup, tensor: torch.Tensor, root: typing.SupportsInt, op: torch._C._distributed_c10d.ReduceOp = <RedOpType.SUM: 0>, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组中的所有进程归约提供的张量。
+
+有关详细信息，请参见 torch.distributed.reduce()。
+
+reduce_scatter(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[torch.Tensor], input_tensors: collections.abc.Sequence[collections.abc.Sequence[torch.Tensor]], opts: torch._C._distributed_c10d.ReduceScatterOptions = <torch._C._distributed_c10d.ReduceScatterOptions object at 0x7f0162ee5cf0>) -> c10d::Work
+
+跨进程组从所有进程中归约并散射输入张量。
+
+有关详细信息，请参见 torch.distributed.reduce_scatter()。
+
+reduce_scatter(self: torch._C._distributed_c10d.ProcessGroup, output: torch.Tensor, input: collections.abc.Sequence[torch.Tensor], op: torch._C._distributed_c10d.ReduceOp = <RedOpType.SUM: 0>, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组从所有进程中归约并散射输入张量。
+
+有关详细信息，请参见 torch.distributed.reduce_scatter()。
+
+跨进程组从所有进程中归约并散射输入张量。
+
+有关详细信息，请参见 torch.distributed.reduce_scatter()。
+
+scatter(self: torch._C._distributed_c10d.ProcessGroup, output_tensors: collections.abc.Sequence[torch.Tensor], input_tensors: collections.abc.Sequence[collections.abc.Sequence[torch.Tensor]], opts: torch._C._distributed_c10d.ScatterOptions = <torch._C._distributed_c10d.ScatterOptions object at 0x7f0162b879f0>) -> c10d::Work
+
+跨进程组从所有进程中 scatter（散射）输入张量。
+
+有关详细信息，请参见 torch.distributed.scatter()。
+
+scatter(self: torch._C._distributed_c10d.ProcessGroup, output_tensor: torch.Tensor, input_tensors: collections.abc.Sequence[torch.Tensor], root: typing.SupportsInt, timeout: datetime.timedelta | None = None) -> c10d::Work
+
+跨进程组从所有进程中 scatter（散射）输入张量。
+
+有关详细信息，请参见 torch.distributed.scatter()。
+
+将张量发送到指定的 rank。
+
+有关详细信息，请参见 torch.distributed.send()。
+
+为所有未来的操作设置默认超时时间。
+
+关闭进程组
+
+获取此进程组的大小。
+
+用于进程组工厂的协议。
+
+获取当前进程组。线程局部方法。
+
+当前进程组。
+
+使用给定的后端和选项创建一个新的进程组。这个组是独立的，不会被全局注册，因此无法通过标准的 torch.distributed.* API 使用。
+
+backend (str) – 用于进程组的后端。
+
+timeout (timedelta) – 集合通信操作的超时时间。
+
+device (Union[str, device]) – 用于进程组的设备。
+
+**kwargs (object) – 所有剩余的参数都将传递给后端构造函数。有关详细信息，请参见特定于后端的文档。
+
+进程组的上下文管理器。线程局部方法。
+
+pg (ProcessGroup) – 要使用的进程组。
+
+Generator[None, None, None]
+
+注册一个新的进程组后端。
+
+name (str) – 后端的名称。
+
+func (ProcessGroupFactory) – 用于创建进程组的函数。
+
+---
+
+## torch.distributed.fsdp.fully_shard#
+
+**URL:** https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html
+
+**目录：**
+- torch.distributed.fsdp.fully_shard#
+- PyTorch FSDP2 (fully_shard)#
+
+创建时间：2024年12月4日 | 最后更新时间：2025年6月16日
+
+PyTorch FSDP2 (RFC) 提供了一个完全分片数据并行（FSDP）的实现，旨在实现高性能的 eager 模式，同时使用逐参数分片以提高可用性
+
+有关更多信息，请参阅 FSDP2 入门教程。
+
+如果您当前正在使用 FSDP1，请考虑使用我们的迁移指南迁移到 FSDP2。
+
+fully_shard(model) 的用户契约如下
+
+对于模型初始化，fully_shard 原地将 model.parameters() 从普通的 torch.Tensor 转换为 DTensor。参数将根据设备网格移动到相应的设备上。
+
+在前向和反向传递之前，前向/反向前置钩子负责 all-gather 参数，并将 model.parameters() 从 DTensor 转换为普通的 torch.Tensor。
+
+在前向和反向传递之后，前向/反向后置钩子释放未分片的参数（不需要通信），并将 model.parameters() 从普通的 torch.Tensor 转换回 DTensor。
+
+对于优化器，它必须使用 DTensor 格式的 model.parameters() 进行初始化，并且优化器步骤应在 DTensor 参数上执行。
+
+请调用 model(input) 而不是 model.forward(input)，以触发前向前置钩子来 all-gather 参数。要使 model.forward(input) 正常工作，用户必须显式调用 model.unshard() 或使用 register_fsdp_forward_method(model, "forward") 来注册该前向方法以进行挂钩。
+
+fully_shard 将参数分组在一起以便进行单一的 all-gather。用户应以自下而上的方式应用 fully_shard。例如，在 Transformer 模型中，应将 fully_shard 应用于根模型之前的每一层。当应用于根模型时，fully_shard 会从每一层中排除 model.parameters()，并将剩余的参数（例如 embeddings、输出投影）分组到一个单一的 all-gather 组中。
+
+type(model) 会原位与 FSDPModule 进行“联合”。例如，如果 model 最初是 nn.Linear 类型，则 fully_shard 会将 type(model) 从 nn.Linear 原位更改为 FSDPLinear。FSDPLinear 既是 nn.Linear 又是 FSDPModule 的实例。它保留了 nn.Linear 的所有方法，同时在 FSDPModule 下公开特定于 FSDP2 的 API，例如 reshard() 和 unshard()。
+
+参数的全限定名保持不变。如果我们调用 model.state_dict()，则在应用 fully_shard 之前和之后 FQN 是相同的。这是因为 fully_shard 没有包装模块，而只是将钩子注册到原始模块。
+
+与 PyTorch FSDP1 (FullyShardedDataParallel) 相比：
+
+与 FSDP1 的扁平参数分片相比，FSDP2 使用基于 DTensor 的 dim-0 逐参数分片以获得更简单的分片表示，同时保持相似的吞吐量性能。更具体地说，FSDP2 在数据并行工作进程之间对每个参数的 dim-0 进行分块（使用 torch.chunk(dim=0)），而 FSDP1 将一组张量展平、拼接后再进行分块，这使得推断每个工作进程上存在哪些数据以及重新分片为不同的并行策略变得复杂。逐参数分片提供了更直观的用户体验，放宽了对冻结参数的限制，并允许无通信的（分片）状态字典，而在 FSDP1 中这通常需要 all-gather。
+
+FSDP2 实现了一种不同的内存管理方法来处理多流使用场景，避免了 torch.Tensor.record_stream。这确保了确定性和预期的内存使用情况，并且不需要像 FSDP1 的 limit_all_gathers=True 那样阻塞 CPU。
+
+FSDP2 公开了用于手动控制预取和集合通信调度的 API，允许高级用户进行更多自定义。有关详细信息，请参见下文中的 FSDPModule 上的方法。
+
+FSDP2 简化了一些 API 表面：例如，FSDP2 不直接支持完整状态字典。相反，用户可以使用 DTensor API（如 DTensor.full_tensor()）或使用更高级别的 API（如 PyTorch 分布式检查点的分布式状态字典 API）自行将包含 DTensor 的分片状态字典重新分片为完整状态字典。此外，一些其他参数已被移除；有关详细信息，请参见此处。
+
+可以在模块上调用前端 API fully_shard：
+
+将完全分片数据并行（FSDP）应用于模块，其中 FSDP 将模块参数、梯度和优化器状态跨数据并行工作进程进行分片，以通过通信成本换取内存的节省。
+
+在初始化时，FSDP 根据 mesh 给定的数据并行工作进程对模块的参数进行分片。在前向传递之前，FSDP 跨数据并行工作进程 all-gather 已分片的参数，以获取用于前向计算的未分片参数。如果 reshard_after_forward 为 True，则 FSDP 在前向传递之后释放未分片的参数，并在计算梯度之前的反向传递中重新 all-gather 它们。在计算梯度之后，FSDP 释放未分片的参数，并跨数据并行工作进程对未分片的梯度进行 reduce-scatter（归约-散射）。
+
+此实现将在 dim-0 上分片的 DTensor 表示为分片参数，而未分片的参数将与模块上的原始参数相同（例如，如果原来是 torch.Tensor，则为 torch.Tensor）。模块上的前向前置钩子会 all-gather 参数，而模块上的前向钩子（如果需要）会释放它们。类似的后向钩子会 all-gather 参数，随后释放参数并 reduce-scatter 梯度。
+
+由于将多个张量组合在一起进行一次集合通信对于通信效率至关重要，因此此实现使这种组合成为一等公民。在模块上调用 fully_shard() 会构造一个组，其中包含 module.parameters() 中的参数，但之前在对子模块的调用中已分配给组的参数除外。这意味着应在模型上自下而上地调用 fully_shard()。每个组的参数将在一次集合通信中被 all-gather，并且其梯度将在一次集合通信中被 reduce-scatter。将模型划分为多个组（“逐层”）可以实现峰值内存节省以及通信与计算的重叠。用户通常不应该只在最顶层的根模块上调用 fully_shard()。
+
+module (Union[nn.Module, List[nn.Module]) – 要使用 FSDP 进行分片并组合在一起以进行通信的模块或模块列表。
+
+mesh (可选[DeviceMesh]) – 此数据并行网格定义了分片和设备。如果是 1D，则参数在 1D 网格上以 (Shard(0),) 放置进行完全分片（FSDP）。如果是 2D，则参数在第 1 维上进行分片，并在第 0 维上复制（HSDP），以 (Replicate(), Shard(0)) 进行放置。网格的设备类型给出了用于通信的设备类型；如果是 CUDA 或类 CUDA 设备类型，我们使用当前设备。
+
+reshard_after_forward (可选[Union[bool, int]]) – 这控制了前向传递后的参数行为，并且可以在内存和通信之间进行权衡：如果为 True，则这会在前向传递之后重新分片参数，并在反向传递中重新 all-gather。如果为 False，则这会在前向传递之后将未分片的参数保留在内存中，并避免在反向传递中执行 all-gather。为了获得最佳性能，我们通常为根模块设置为 False，因为在反向传递开始时通常立即需要根模块。如果为 None，则对于非根模块设置为 True，对于根模块设置为 False。如果是整数，则表示前向传递后要重新分片到的 world size。它应该是 mesh 分片维度大小的非平凡因子（即不包括 1 和维度大小本身）。选择之一可以是节点内大小（例如 torch.cuda.device_count()）。这允许在反向传递中以较小的 world size 执行 all-gather，代价是比设置为 True 占用更高的内存。在前向传递之后，注册到模块的参数取决于此设置：如果为 True，注册的参数是分片参数；如果为 False，则是未分片的参数；否则为重新分片到较小网格的参数。要修改前向和反向传递之间的参数，注册的参数必须是分片参数。对于 False 或整数，这可以通过 reshard() 手动重新分片来完成。
+
+这控制了前向传递后的参数行为，并且可以在内存和通信之间进行权衡：
+
+如果为 True，则这会在前向传递之后重新分片参数，并在反向传递中重新 all-gather。
+
+如果为 False，则这会在前向传递之后将未分片的参数保留在内存中，并避免在反向传递中执行 all-gather。为了获得最佳性能，我们通常为根模块设置为 False，因为在反向传递开始时通常立即需要根模块。
+
+如果为 None，则对于非根模块设置为 True，对于根模块设置为 False。
+
+如果是整数，则表示前向传递后要重新分片到的 world size。它应该是 mesh 分片维度大小的非平凡因子（即不包括 1 和维度大小本身）。选择之一可以是节点内大小（例如 torch.cuda.device_count()）。这允许在反向传递中以较小的 world size 执行 all-gather，代价是比设置为 True 占用更高的内存。
+
+在前向传递之后，注册到模块的参数取决于此设置：如果为 True，注册的参数是分片参数；如果为 False，则是未分片的参数；否则为重新分片到较小网格的参数。要修改前向和反向传递之间的参数，注册的参数必须是分片参数。对于 False 或整数，这可以通过 reshard() 手动重新分片来完成。
+
+shard_placement_fn (可选[Callable[[nn.Parameter], 可选[Shard]]]) – 此可调用对象可用于覆盖参数的分片放置，以便在 dim-0 以外的维度上对参数进行分片。如果此可调用对象返回了一个 Shard 放置（非 None），则 FSDP 将根据该放置进行分片（例如 Shard(1)）。如果在非零维度上进行分片，我们目前要求均匀分片，即该维度上的张量维度大小必须能被 FSDP 分片 mesh 大小整除。
+
+mp_policy (MixedPrecisionPolicy) – 这控制了混合精度策略，它为此模块提供参数/归约混合精度。有关详细信息，请参见 MixedPrecisionPolicy。
+
+offload_policy (OffloadPolicy) – 这控制了卸载策略，它提供了参数/梯度/优化器状态的卸载。有关详细信息，请参见 OffloadPolicy 及其子类。
+
+ignored_params (可选[set[nn.Parameter]]) – 可选(Set[nn.Parameter]): FSDP 将忽略的参数集。它们不会被分片，也不会在初始化期间移动到设备上，也不会在反向传递中对其梯度进行归约。
+
+应用了 FSDP 的模块（原位操作）。
+
+重新分片模块的参数，如果未分片的参数已分配则释放它们，并将分片参数注册到模块。此方法不是递归的。
+
+hook (Callable[[torch.Tensor], None]) – 用户自定义的 all-reduce 钩子，预期签名为 hook(reduce_output: torch.Tensor) -> None，其中如果仅使用 FSDP，reduce_output 为 reduce-scatter 输出；如果使用原生 HSDP，则为 all-reduce 输出。
+
+stream (可选[torch.cuda.Stream]) – 运行 all-reduce 钩子的流。仅在不使用原生 HSDP 时才应设置此项。如果使用原生 HSDP，该钩子将在原生 HSDP all-reduce 内部使用的定义的 all-reduce 流中运行。
+
+设置用于通过集合通信发送和接收数据的临时暂存缓冲区是否应使用 ProcessGroup 本身提供的自定义优化分配器进行分配（如果有）。这可能允许 ProcessGroup 更高效。例如，当使用 NCCL 时，这使其能够通过 SHARP（对于 NVLink 和/或 InfiniBand）利用零拷贝传输。
+
+这不能与 set_custom_all_gather() 或 set_custom_reduce_scatter() 一起使用，因为这些 API 允许对每次通信进行更细粒度的控制，而此方法无法确定它们的暂存缓冲区分配策略。
+
+enable (bool) – 是否开启 ProcessGroup 分配。
+
+覆盖默认的 all_gather 通信行为，以更好地控制通信和内存使用。有关详细信息，请参见 Comm 和 ReduceScatter。
+
+comm (AllGather) – 自定义 all-gather 通信。
+
+覆盖默认的 reduce_scatter 通信行为，以更好地控制通信和内存使用。有关详细信息，请参见 Comm 和 ReduceScatter。
+
+comm (ReduceScatter) – 自定义 reduce_scatter 通信。
+
+设置是否要求底层集合通信原语专门使用“求和”类型的归约，即使这会带来单独的额外预缩放或后缩放操作代价。例如，需要这样做是因为 NCCL 目前仅支持此类集合通信的零拷贝传输。
+
+注意：对于 MTIA 设备，这始终是隐式启用的。
+
+注意：如果在 FSDP 设置下使用了 set_all_reduce_hook，调用者需要确保跨 FSDP 单元的自定义 all-reduce 也遵循此策略，因为 FSDP 无法再自动处理该情况。
+
+enable (bool) – 是否仅使用 ReduceOp.SUM 进行通信。
+
+为梯度归约设置自定义除法因子。这可能会使用 NCCL 的 PreMulSum 使用自定义归约操作，允许在归约之前乘以该因子。
+
+factor (float) – 自定义
+
+## 分布式通信包 - torch.distributed#
+
+**URL:** https://pytorch.org/docs/stable/distributed.html
+
+**目录:**
+- 分布式通信包 - torch.distributed#
+- 后端#
+  - PyTorch 自带的后端#
+  - 使用哪个后端？#
+  - 常用环境变量#
+    - 选择要使用的网络接口#
+    - 其他 NCCL 环境变量#
+- 基础#
+- 初始化#
+  - TCP 初始化#
+
+创建时间：2017 年 7 月 12 日 | 最后更新时间：2025 年 9 月 4 日
+
+有关与分布式训练相关的所有功能的简要介绍，请参阅 PyTorch Distributed 概览。
+
+torch.distributed 支持四种内置后端，每种后端具有不同的功能。下表显示了每种后端在 CPU 或 GPU 上可用的功能。对于 NCCL，GPU 指的是 CUDA GPU，而对于 XCCL 指的是 XPU GPU。
+
+只有用于构建 PyTorch 的 MPI 实现支持 CUDA 时，MPI 才支持 CUDA。
+
+PyTorch 分布式包支持 Linux（稳定版）、MacOS（稳定版）和 Windows（原型版）。在 Linux 上默认情况下，Gloo 和 NCCL 后端会被构建并包含在 PyTorch 分布式中（NCCL 仅在使用 CUDA 构建时包含）。MPI 是一个可选后端，只有在从源代码构建 PyTorch 时才能包含它。（例如，在安装了 MPI 的主机上构建 PyTorch。）
+
+从 PyTorch v1.8 开始，Windows 支持除 NCCL 外的所有集合通信后端，如果 init_process_group() 的 init_method 参数指向一个文件，它必须遵循以下模式：
+
+本地文件系统，init_method="file:///d:/tmp/some_file"
+
+共享文件系统，init_method="file://////{machine_name}/{share_folder_name}/some_file"
+
+与 Linux 平台上一样，您可以通过设置环境变量 MASTER_ADDR 和 MASTER_PORT 来启用 TcpStore。
+
+过去，我们经常被问到：“我应该使用哪个后端？”。
+
+在进行使用 CUDA GPU 的分布式训练时使用 NCCL 后端。
+
+在进行使用 XPU GPU 的分布式训练时使用 XCCL 后端。
+
+在进行使用 CPU 的分布式训练时使用 Gloo 后端。
+
+带有 InfiniBand 互连的 GPU 主机
+
+使用 NCCL，因为它是目前唯一支持 InfiniBand 和 GPUDirect 的后端。
+
+带有以太网互连的 GPU 主机
+
+使用 NCCL，因为它目前为分布式 GPU 训练提供了最佳性能，特别是对于多进程单节点或多节点分布式训练。如果您在使用 NCCL 时遇到任何问题，请使用 Gloo 作为备用选项。（请注意，目前 Gloo 在 GPU 上的运行速度比 NCCL 慢。）
+
+带有 InfiniBand 互连的 CPU 主机
+
+如果您的 InfiniBand 启用了 IP over IB，请使用 Gloo，否则请使用 MPI。我们计划在即将发布的版本中为 Gloo 添加 InfiniBand 支持。
+
+带有以太网互连的 CPU 主机
+
+使用 Gloo，除非您有特定原因要使用 MPI。
+
+默认情况下，NCCL 和 Gloo 后端都会尝试寻找合适的网络接口来使用。如果自动检测到的接口不正确，您可以使用以下环境变量（适用于各自的后端）进行覆盖：
+
+NCCL_SOCKET_IFNAME，例如 export NCCL_SOCKET_IFNAME=eth0
+
+GLOO_SOCKET_IFNAME，例如 export GLOO_SOCKET_IFNAME=eth0
+
+如果您使用的是 Gloo 后端，您可以通过用逗号分隔来指定多个接口，如下所示：export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3。后端将以轮询的方式在这些接口上分配操作。所有进程必须在此变量中指定相同数量的接口，这一点至关重要。
+
+调试 - 如果 NCCL 发生故障，您可以设置 NCCL_DEBUG=INFO 来打印明确的警告信息以及基本的 NCCL 初始化信息。
+
+您还可以使用 NCCL_DEBUG_SUBSYS 获取有关 NCCL 特定方面的更多详细信息。例如，NCCL_DEBUG_SUBSYS=COLL 将打印集合调用的日志，这在调试挂起（尤其是由集合类型或消息大小不匹配引起的挂起）时非常有用。如果拓扑检测失败，设置 NCCL_DEBUG_SUBSYS=GRAPH 来检查详细的检测结果，并在需要 NCCL 团队提供进一步帮助时将其保存为参考会很有帮助。
+
+性能调优 - NCCL 根据其拓扑检测执行自动调优，以节省用户的调优精力。在某些基于 socket 的系统上，用户仍然可以尝试调整 NCCL_SOCKET_NTHREADS 和 NCCL_NSOCKS_PERTHREAD 来增加 socket 网络带宽。NCCL 已为一些云提供商（如 AWS 或 GCP）预调优了这两个环境变量。
+
+有关 NCCL 环境变量的完整列表，请参阅 NVIDIA NCCL 的官方文档。
+
+你可以使用 torch.distributed.ProcessGroupNCCL.NCCLConfig 和 torch.distributed.ProcessGroupNCCL.Options 进一步 tune NCCL 通信器。在解释器中使用 help（例如 help(torch.distributed.ProcessGroupNCCL.NCCLConfig)）来了解更多关于它们的信息。
+
+torch.distributed 包为在一台或多台机器上运行的多个计算节点间的多进程并行提供了 PyTorch 支持和通信原语。torch.nn.parallel.DistributedDataParallel() 类在此功能的基础上构建，作为任何 PyTorch 模型的包装器，提供同步分布式训练。这与 Multiprocessing 包 - torch.multiprocessing 和 torch.nn.DataParallel() 提供的并行类型不同，它支持多台网络连接的机器，并且用户必须为每个进程显式启动主训练脚本的单独副本。
+
+在单机同步的情况下，torch.distributed 或 torch.nn.parallel.DistributedDataParallel() 包装器在数据并行方面可能仍然比其他方法（包括 torch.nn.DataParallel()）具有优势：
+
+每个进程维护自己的优化器，并在每次迭代中执行完整的优化步骤。虽然这看起来可能是冗余的，因为梯度已经跨进程收集并求平均，因此对于每个进程来说都是相同的，但这意味着不需要参数广播步骤，从而减少了在节点之间传输张量所花费的时间。
+
+每个进程都包含一个独立的 Python 解释器，从而消除了从单个 Python 进程驱动多个执行线程、模型副本或 GPU 所带来的额外解释器开销和“GIL 抖动”。这对于大量使用 Python 运行时的模型（包括具有循环层或许多小组件的模型）尤为重要。
+
+在调用任何其他方法之前，需要使用 torch.distributed.init_process_group() 或 torch.distributed.device_mesh.init_device_mesh() 函数初始化此包。两者都会阻塞，直到所有进程都加入。
+
+初始化不是线程安全的。进程组的创建应该从单个线程执行，以防止跨 rank 的“UUID”分配不一致，并防止初始化期间可能导致挂起的竞争。
+
+如果分布式包可用，则返回 True。
+
+否则，torch.distributed 不会暴露任何其他 API。目前，torch.distributed 可在 Linux、MacOS 和 Windows 上使用。从源代码构建 PyTorch 时，设置 USE_DISTRIBUTED=1 以启用它。目前，Linux 和 Windows 的默认值为 USE_DISTRIBUTED=1，MacOS 的默认值为 USE_DISTRIBUTED=0。
+
+初始化默认分布式进程组。
+
+这也将初始化分布式包。
+
+显式指定 store、rank 和 world_size。
+
+指定 init_method（一个 URL 字符串），指示在哪里/如何发现对等节点。可选择指定 rank 和 world_size，或者在 URL 中编码所有必需的参数并省略它们。
+
+如果两者都未指定，则假定 init_method 为“env://”。
+
+backend (str 或 Backend, 可选) – 要使用的后端。根据构建时的配置，有效值包括 mpi、gloo、nccl、ucc、xccl 或由第三方插件注册的后端。从 2.6 版本开始，如果未提供 backend，c10d 将使用为 device_id 关键字参数（如果提供）指示的设备类型注册的后端。目前已知的默认注册是：cuda 对应 nccl，cpu 对应 gloo，xpu 对应 xccl。如果未提供 backend 和 device_id，c10d 将检测运行时机器上的加速器，并使用为检测到的加速器（或 cpu）注册的后端。此字段可以作为小写字符串给出（例如，"gloo"），也可以通过 Backend 属性访问（例如，Backend.GLOO）。如果在每台机器上使用具有 nccl 后端的多个进程，则每个进程必须对其使用的每个 GPU 具有独占访问权，因为在进程之间共享 GPU 可能导致死锁或 NCCL 无效使用。ucc 后端是实验性的。可以通过 get_default_backend_for_device() 查询设备的默认 backend。
+
+init_method (str, 可选) – 指定如何初始化进程组的 URL。如果没有指定 init_method 或 store，则默认为“env://”。与 store 互斥。
+
+world_size (int, 可选) – 参与作业的进程数。如果指定了 store 则为必需。
+
+rank (int, 可选) – 当前进程的 rank（它应该是 0 和 world_size-1 之间的数字）。如果指定了 store 则为必需。
+
+store (Store, 可选) – 所有工作进程均可访问的键/值存储，用于交换连接/地址信息。与 init_method 互斥。
+
+timeout (timedelta, 可选) – 针对进程组执行的操作的超时时间。NCCL 的默认值为 10 分钟，其他后端为 30 分钟。在此持续时间之后，集合通信将被异步中止，进程将崩溃。这样做是因为 CUDA 执行是异步的，并且由于失败的异步 NCCL 操作可能导致后续 CUDA 操作在损坏的数据上运行，因此继续执行用户代码不再安全。设置了 TORCH_NCCL_BLOCKING_WAIT 时，进程将阻塞并等待此超时。
+
+group_name (str, 可选, 已弃用) – 组名。此参数被忽略
+
+pg_options (ProcessGroupOptions, 可选) – 进程组选项，指定在构造特定进程组期间需要传入哪些附加选项。目前，我们支持的唯一选项是用于 nccl 后端的 ProcessGroupNCCL.Options，可以指定 is_high_priority_stream，以便在有计算内核等待时，nccl 后端可以选取高优先级的 CUDA 流。有关配置 nccl 的其他可用选项，请参阅 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t
+
+device_id (torch.device | int, 可选) – 此进程将在其上运行的单个特定设备，允许进行特定于后端的优化。目前这有两个效果，仅在 NCCL 下：通信器立即形成（立即调用 ncclCommInit* 而不是正常的惰性调用），并且子组将在可能的情况下使用 ncclCommSplit 以避免组创建的不必要开销。如果您想及早了解 NCCL 初始化错误，也可以使用此字段。如果提供的是 int，则 API 假定将使用编译时的加速器类型。
+
+要启用 backend == Backend.MPI，需要在支持 MPI 的系统上从源代码构建 PyTorch。
+
+对多个后端的支持是实验性的。目前，当未指定后端时，将同时创建 gloo 和 nccl 后端。gloo 后端将用于具有 CPU 张量的集合通信，而 nccl 后端将用于具有 CUDA 张量的集合通信。可以通过传入格式为“<device_type>:<backend_name>,<device_type>:<backend_name>”的字符串（例如“cpu:gloo,cuda:custom_backend”）来指定自定义后端。
+
+根据 device_type、mesh_shape 和 mesh_dim_names 参数初始化 DeviceMesh。
+
+这将创建一个具有 n 维数组布局的 DeviceMesh，其中 n 是 mesh_shape 的长度。如果提供了 mesh_dim_names，则每个维度标记为 mesh_dim_names[i]。
+
+init_device_mesh 遵循 SPMD 编程模型，这意味着相同的 PyTorch Python 程序在集群中的所有进程/rank 上运行。确保 mesh_shape（描述设备布局的 nD 数组的维度）在所有 rank 上都是相同的。不一致的 mesh_shape 可能会导致挂起。
+
+如果未找到进程组，init_device_mesh 将在后台初始化分布式通信所需的分布式进程组或组群。
+
+device_type (str) – mesh 的设备类型。目前支持：“cpu”、“cuda/cuda-like”、“xpu”。不允许传入带有 GPU 索引的设备类型，例如“cuda:0”。
+
+mesh_shape (Tuple[int]) – 定义描述设备布局的多维数组维度的元组。
+
+mesh_dim_names (Tuple[str], 可选) – 分配给描述设备布局的多维数组每个维度的 mesh 维度名称元组。其长度必须与 mesh_shape 的长度相匹配。mesh_dim_names 中的每个字符串必须是唯一的。
+
+backend_override (Dict[int | str, tuple[str, Options] | str | Options], 可选) – 对将为每个 mesh 维度创建的部分或全部 ProcessGroup 的覆盖。每个键可以是维度的索引或其名称（如果提供了 mesh_dim_names）。每个值可以是一个包含后端名称及其选项的元组，也可以只是这两个组件中的一个（在这种情况下，另一个将被设置为其默认值）。
+
+一个代表设备布局的 DeviceMesh 对象。
+
+检查默认进程组是否已初始化。
+
+检查 MPI 后端是否可用。
+
+检查 NCCL 后端是否可用。
+
+检查 Gloo 后端是否可用。
+
+检查 XCCL 后端是否可用。
+
+检查此进程是否由 torch.distributed.elastic（又名 torchelastic）启动。
+
+TORCHELASTIC_RUN_ID 环境变量的存在被用作确定当前进程是否由 torchelastic 启动的代理。这是一个合理的代理，因为 TORCHELASTIC_RUN_ID 映射到 rendezvous id，该 id 始终为非空值，指示用于对等发现目的的作业 ID。
+
+返回给定设备的默认后端。
+
+device (Union[str, torch.device]) – 要获取其默认后端的设备。
+
+给定设备的默认后端，以小写字符串形式返回。
+
+目前支持三种初始化方法：
+
+使用 TCP 初始化有两种方法，都需要一个所有进程都可以访问的网络地址和所需的 world_size。第一种方法需要指定一个属于 rank 0 进程的地址。这种初始化方法要求所有进程都手动指定 rank。
+
+请注意，最新的分布式包中不再支持多播地址。group_name 也已弃用。
+
+另一种初始化方法利用了一个组中所有机器都可以看到并共享的文件系统，以及所需的 world_size。URL 应以 file:// 开头，并包含共享文件系统上一个（位于现有目录中的）不存在文件的路径。如果文件不存在，文件系统初始化将自动创建该文件，但不会删除该文件。因此，您有责任确保在下一次对相同文件路径/名称调用 init_process_group() 之前清理该文件。
+
+请注意，最新的分布式包中不再支持自动 rank 分配，并且 group_name 也已弃用。
+
+此方法假定文件系统支持使用 fcntl 进行锁定 - 大多数本地系统和 NFS 都支持它。
+
+此方法将始终创建文件，并尽最大努力在程序结束时清理和删除该文件。换句话说，每次使用 file init 方法进行初始化时都需要一个全新的空文件，以便初始化成功。如果再次使用先前初始化使用过的文件（碰巧未被清理），这是意外行为，通常会导致死锁和失败。因此，即使此方法会尽最大努力清理文件，如果自动删除碰巧未成功，您有责任确保在训练结束时删除该文件，以防止下次再次重复使用相同的文件。如果您打算在同一文件名上多次调用 init_process_group()，这一点尤为重要。换句话说，如果文件未被删除/清理，而您再次对该文件调用 init_process_group()，预计会发生失败。这里的经验法则是，确保每次调用 init_process_group() 时文件都不存在或为空。
+
+此方法将从环境变量中读取配置，允许完全自定义获取信息的方式。要设置的变量有：
+
+MASTER_PORT - 必填；必须是 rank 0 机器上的一个空闲端口
+
+MASTER_ADDR - 必填（rank 0 除外）；rank 0 节点的地址
+
+WORLD_SIZE - 必填；可以在此处设置，或在调用 init 函数时设置
+
+RANK - 必填；可以在此处设置，或在调用 init 函数时设置
+
+rank 为 0 的机器将用于建立所有连接。
+
+这是默认方法，意味着不必指定 init_method（或者可以是 env://）。
+
+TORCH_GLOO_LAZY_INIT - 按需建立连接，而不是使用全网状连接，这可以大大改善非 all2all 操作的初始化时间。
+
+一旦运行了 torch.distributed.init_process_group()，就可以使用以下函数。要检查进程组是否已经初始化，请使用 torch.distributed.is_initialized()。
+
+用于后端的类枚举类。
+
+可用后端：GLOO、NCCL、UCC、MPI、XCCL 和其他注册后端。
+
+此类的值为小写字符串，例如 "gloo"。它们可以作为属性访问，例如 Backend.NCCL。
+
+这个类可以被直接调用来解析字符串，例如，Backend(backend_str) 将检查 backend_str 是否有效，如果有效则返回解析后的小写字符串。它也接受大写字符串，例如，Backend("GLOO") 返回 "gloo"。
+
+Backend.UNDEFINED 条目存在，但仅用作某些字段的初始值。用户既不应直接使用它，也不应假定它存在。
+
+使用给定的名称和实例化函数注册一个新的后端。
+
+第三方 ProcessGroup 扩展使用此类的注册新后端。
+
+name (str) – ProcessGroup 扩展的后端名称。它应与 init_process_group() 中的名称相匹配。
+
+func (函数) – 实例化后端的函数处理程序。该函数应在后端扩展中实现，并接受四个参数，包括 store、rank、world_size 和 timeout。
+
+extended_api (bool, 可选) – 后端是否支持扩展参数结构。默认: False。如果设置为 True，后端将获得一个 c10d::DistributedBackendOptions 实例，以及一个由后端实现定义的进程组选项对象。
+
+device (str 或 str 列表, 可选) – 此后端支持的设备类型，例如 “cpu”、“cuda” 等。如果为 None，则假定同时支持“cpu”和“cuda”
+
+这种对第三方后端的支持是实验性的，可能会发生变化。
+
+返回给定进程组的后端。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。默认是通用的主进程组。如果指定了另一个特定的组，则调用进程必须是该组的一部分。
+
+给定进程组的后端，以小写字符串形式返回。
+
+返回当前进程在提供的组中的 rank，否则为默认组。
+
+Rank 是分配给分布式进程组内每个进程的唯一标识符。它们始终是范围从 0 到 world_size 的连续整数。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+进程组的 rank，如果不属于该组，则为 -1
+
+返回当前进程组中的进程数。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+进程组的 world size，如果不属于该组，则为 -1
+
+在退出时通过调用 destroy_process_group() 清理资源非常重要。
+
+最简单的模式是在训练脚本中不再需要通信的时候（通常在 main() 结尾附近），通过调用 destroy_process_group() 并为 group 参数传入默认值 None 来销毁每个进程组和后端。该调用应在每个训练器进程上执行一次，而不是在外部进程启动器级别执行。
+
+如果 pg 中的所有 rank 在超时时间内未调用 destroy_process_group()，尤其是当应用程序中有多个进程组时（例如用于 N-D 并行），在退出时可能会挂起。这是因为 ProcessGroupNCCL 的析构函数会调用 ncclCommAbort，而它必须被共同调用，但是如果由 python 的 GC 调用，调用 ProcessGroupNCCL 的析构函数的顺序是不确定的。调用 destroy_process_group() 有助于确保在所有 rank 之间以一致的顺序调用 ncclCommAbort，并避免在 ProcessGroupNCCL 的析构函数执行期间调用 ncclCommAbort。
+
+destroy_process_group 也可用于销毁单个进程组。一个用例可能是容错训练，其中进程组可能在运行期间被销毁，然后初始化一个新的进程组。在这种情况下，_在_调用 destroy 之后且随后初始化之前，使用 torch.distributed 原语以外的一些手段同步训练器进程是至关重要的。由于实现这种同步很困难，这种行为目前是不受支持/未经过测试的，并且被认为是一个已知问题。如果这个用例阻碍了您，请提交 github issue 或 RFC。
+
+默认情况下，集合通信在默认组（也称为 world）上运行，并要求所有进程都进入分布式函数调用。然而，一些工作负载可以从更细粒度的通信中获益。这就是分布式组发挥作用的地方。new_group() 函数可用于创建包含所有进程任意子集的新组。它返回一个不透明的组句柄，可以作为 group 参数提供给所有集合通信（集合通信是在某些众所周知的编程模式中用于交换信息的分布式函数）。
+
+创建一个新的分布式组。
+
+此函数要求主组中的所有进程（即作为分布式作业一部分的所有进程）都进入此函数，即使它们不打算成为该组的成员。此外，在所有进程中，组应该以相同的顺序创建。
+
+安全的并发使用：当使用具有 NCCL 后端的多个进程组时，用户必须确保跨 rank 的集合通信具有全局一致的执行顺序。
+
+如果进程内的多个线程发出集合通信，则必须进行显式同步以确保一致的顺序。
+
+当使用 torch.distributed 通信 API 的异步变体时，将返回一个 work 对象，并且通信内核会在单独的 CUDA 流上排队，从而允许通信和计算重叠。一旦在一个进程组上发出了一个或多个异步操作，在使用另一个进程组之前，必须通过调用 work.wait() 将它们与其他 CUDA 流同步。
+
+有关更多详细信息，请参阅并发使用多个 NCCL 通信器 <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently>。
+
+ranks (list[int]) – 组成员 rank 的列表。如果为 None，将被设置为所有 rank。默认为 None。
+
+timeout (timedelta, 可选) – 有关详细信息和默认值，请参见 init_process_group。
+
+backend (str 或 Backend, 可选) – 要使用的后端。根据构建时的配置，有效值为 gloo 和 nccl。默认使用与全局组相同的后端。此字段应以小写字符串形式给出（例如 "gloo"），也可以通过 Backend 属性访问（例如 Backend.GLOO）。如果传入 None，将使用与默认进程组对应的后端。默认为 None。
+
+pg_options (ProcessGroupOptions, 可选) – 进程组选项，指定在构造特定进程组期间需要传入哪些附加选项。即对于 nccl 后端，可以指定 is_high_priority_stream，以便进程组可以选取高优先级的 CUDA 流。有关配置 nccl 的其他可用选项，请参阅 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-tuse_local_synchronization (bool, 可选): 在进程组创建结束时执行组本地屏障。这不同之处在于，非成员 rank 不需要调用 API 也不加入屏障。
+
+group_desc (str, 可选) – 描述进程组的字符串。
+
+device_id (torch.device, 可选) – 将此进程“绑定”到的单个特定设备，如果提供此字段，new_group 调用将尝试立即为该设备初始化通信后端。
+
+一个分布式组的句柄，可提供给集合调用，或者如果该 rank 不属于 ranks，则为 GroupMember.NON_GROUP_MEMBER。
+
+注意 use_local_synchronization 无法与 MPI 一起使用。
+
+注意 虽然在集群较大且进程组较小的情况下，use_local_synchronization=True 可以显著提高速度，但必须小心，因为它改变了集群行为，因为非成员 rank 不加入组的 barrier()。
+
+注意 当每个 rank 创建多个重叠的进程组时，use_local_synchronization=True 可能导致死锁。为避免这种情况，请确保所有 rank 遵循相同的全局创建顺序。
+
+将全局 rank 转换为组 rank。
+
+global_rank 必须是该组的一部分，否则会引发 RuntimeError。
+
+group (ProcessGroup) – 要查找相对 rank 的 ProcessGroup。
+
+global_rank (int) – 要查询的全局 rank。
+
+global_rank 相对于组的组 rank
+
+注意 在默认进程组上调用此函数返回原值
+
+将组 rank 转换为全局 rank。
+
+group_rank 必须是该组的一部分，否则会引发 RuntimeError。
+
+group (ProcessGroup) – 要从中查找全局 rank 的 ProcessGroup。
+
+group_rank (int) – 要查询的组 rank。
+
+group_rank 相对于组的全局 rank
+
+注意 在默认进程组上调用此函数返回原值
+
+获取与组关联的所有 rank。
+
+group (可选[ProcessGroup]) – 要从中获取所有 rank 的 ProcessGroup。如果为 None，则使用默认进程组。
+
+按组 rank 排序的全局 rank 列表。
+
+DeviceMesh 是管理进程组（或 NCCL 通信器）的更高级别的抽象。它允许用户轻松创建节点间和节点内的进程组，而无需担心如何为不同的子进程组正确设置 rank，并且它有助于轻松管理这些分布式进程组。init_device_mesh() 函数可用于创建新的 DeviceMesh，并使用一个描述设备拓扑的 mesh 形状。
+
+DeviceMesh 代表一个设备网格，其中设备的布局可以表示为 n-d 维数组，而 n-d 维数组的每个值是默认进程组 rank 的全局 ID。
+
+DeviceMesh 可用于设置跨集群的 N 维设备连接，并管理用于 N 维并行的 ProcessGroup。通信可以分别发生在 DeviceMesh 的每个维度上。DeviceMesh 尊重用户已经选择的设备（即如果用户在 DeviceMesh 初始化之前调用 torch.cuda.set_device），并且如果用户事先未设置设备，它将为当前进程选择/设置设备。请注意，手动选择设备必须在 DeviceMesh 初始化之前进行。
+
+与 DTensor API 一起使用时，DeviceMesh 也可以用作上下文管理器。
+
+DeviceMesh 遵循 SPMD 编程模型，这意味着相同的 PyTorch Python 程序在集群中的所有进程/rank 上运行。因此，用户需要确保 mesh 数组（描述设备的布局）在所有 rank 上是相同的。不一致的 mesh 会导致静默挂起。
+
+device_type (str) – mesh 的设备类型。目前支持：“cpu”、“cuda/cuda-like”。
+
+mesh (ndarray) – 描述设备布局的多维数组或整数张量，其中 ID 是默认进程组的全局 ID。
+
+代表设备布局的 DeviceMesh 对象。
+
+下面的程序以 SPMD 方式在每个进程/rank 上运行。在这个示例中，我们有 2 台主机，每台主机有 4 个 GPU。对 mesh 第一维的规约将跨列 (0, 4), .. 和 (3, 7) 进行，对 mesh 第二维的规约将跨行 (0, 1, 2, 3) 和 (4, 5, 6, 7) 进行。
+
+从现有的 ProcessGroup 或现有 ProcessGroup 列表中构造具有 device_type 的 DeviceMesh。
+
+构造的设备 mesh 的维度数等于传入的组数。例如，如果传入单个进程组，则生成的 DeviceMesh 是一维 mesh。如果传入包含 2 个进程组的列表，则生成的 DeviceMesh 是二维 mesh。
+
+如果传入了多于一个组，则必须提供 mesh 和 mesh_dim_names 参数。传入的进程组的顺序决定了 mesh 的拓扑。例如，第一个进程组将是 DeviceMesh 的第 0 维。传入的 mesh 张量必须与传入的进程组数量具有相同的维数，并且 mesh 张量中的维度顺序必须与传入的进程组的顺序相匹配。
+
+group (ProcessGroup 或 list[ProcessGroup]) – 现有的 ProcessGroup 或现有的 ProcessGroup 列表。
+
+device_type (str) – mesh 的设备类型。目前支持：“cpu”、“cuda/cuda-like”。不允许传入带有 GPU 索引的设备类型，例如“cuda:0”。
+
+mesh (torch.Tensor 或 ArrayLike, 可选) – 描述设备布局的多维数组或整数张量，其中 ID 是默认进程组的全局 ID。默认为 None。
+
+mesh_dim_names (tuple[str], 可选) – 分配给描述设备布局的多维数组每个维度的 mesh 维度名称元组。其长度必须与 mesh_shape 的长度相匹配。mesh_dim_names 中的每个字符串必须是唯一的。默认为 None。
+
+代表设备布局的 DeviceMesh 对象。
+
+返回所有 mesh 维度的 ProcessGroup 列表。
+
+一个 ProcessGroup 对象列表。
+
+list[torch.distributed.distributed_c10d.ProcessGroup]
+
+返回此 rank 相对于 mesh 所有维度的相对索引。如果此 rank 不是 mesh 的一部分，则返回 None。
+
+返回由 mesh_dim 指定的单个 ProcessGroup，或者如果未指定 mesh_dim 并且 DeviceMesh 是一维的，则返回 mesh 中唯一的 ProcessGroup。
+
+mesh_dim (str/python:int, 可选) – 它可以是 mesh 维度的名称或索引
+
+None。 (mesh 维度的。默认是) –
+
+一个 ProcessGroup 对象。
+
+返回 DeviceMesh 的给定 mesh_dim 的本地 rank。
+
+mesh_dim (str/python:int, 可选) – 它可以是 mesh 维度的名称或索引
+
+None。 (mesh 维度的。默认是) –
+
+表示本地 rank 的整数。
+
+下面的程序以 SPMD 方式在每个进程/rank 上运行。在这个示例中，我们有 2 台主机，每台主机有 4 个 GPU。在 rank 0、1、2、3 上调用 mesh_2d.get_local_rank(mesh_dim=0) 将返回 0。在 rank 4、5、6、7 上调用 mesh_2d.get_local_rank(mesh_dim=0) 将返回 1。在 rank 0、4 上调用 mesh_2d.get_local_rank(mesh_dim=1) 将返回 0。在 rank 1、5 上调用 mesh_2d.get_local_rank(mesh_dim=1) 将返回 1。在 rank 2、6 上调用 mesh_2d.get_local_rank(mesh_dim=1) 将返回 2。在 rank 3、7 上调用 mesh_2d.get_local_rank(mesh_dim=1) 将返回 3。
+
+返回当前的全局 rank。
+
+同步发送一个张量。
+
+NCCL 后端不支持 tag。
+
+tensor (Tensor) – 要发送的张量。
+
+dst (int) – 全局进程组上的目标 rank（与 group 参数无关）。目标 rank 不应与当前进程的 rank 相同。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+tag (int, 可选) – 用于将发送与远程接收进行匹配的标签
+
+group_dst (int, 可选) – 组上的目标 rank。不能同时指定 dst 和 group_dst。
+
+同步接收一个张量。
+
+NCCL 后端不支持 tag。
+
+tensor (Tensor) – 要用接收到的数据填充的张量。
+
+src (int, 可选) – 全局进程组上的源 rank（与 group 参数无关）。如果未指定，将从任何进程接收。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+tag (int, 可选) – 用于将接收与远程发送进行匹配的标签
+
+group_src (int, 可选) – 组上的目标 rank。不能同时指定 src 和 group_src。
+
+发送方 rank，如果不属于该组，则为 -1
+
+在使用时，isend() 和 irecv() 返回分布式请求对象。通常，此对象的类型是未指定的，因为永远不应该手动创建它们，但它们保证支持两个方法：
+
+is_completed() - 如果操作已完成，则返回 True
+
+wait() - 将阻塞进程直到操作完成。一旦返回，保证 is_completed() 返回 True。
+
+异步发送一个张量。
+
+在请求完成之前修改张量会导致未定义的行为。
+
+NCCL 后端不支持 tag。
+
+与阻塞式的 send 不同，isend 允许 src == dst rank，即向自己发送。
+
+tensor (Tensor) – 要发送的张量。
+
+dst (int) – 全局进程组上的目标 rank（与 group 参数无关）
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+tag (int, 可选) – 用于将发送与远程接收进行匹配的标签
+
+group_dst (int, 可选) – 组上的目标 rank。不能同时指定 dst 和 group_dst
+
+分布式请求对象。如果不属于该组，则为 None
+
+异步接收一个张量。
+
+NCCL 后端不支持 tag。
+
+与阻塞式的 recv 不同，irecv 允许 src == dst rank，即从自己接收。
+
+tensor (Tensor) – 要用接收到的数据填充的张量。
+
+src (int, 可选) – 全局进程组上的源 rank（与 group 参数无关）。如果未指定，将从任何进程接收。
+
+group (ProcessGroup, 可选) – 要在其上工作的进程组。如果为 None，则使用默认进程组。
+
+tag (int, 可选) – 用于将接收与远程发送进行匹配的标签
+
+group_src (int, 可选) – 组上的目标 rank。不能同时指定 src 和 group_src。
+
+分布式请求对象。如果不属于该组，则为 None
+
+同步发送 object_list 中的可序列化对象。
+
+类似于 send()，但可以传入 Python 对象。请注意，object_list 中的所有对象都必须是可序列化的才能被发送。
+
+object_list (List[Any]) – 要发送的输入对象列表。每个对象都必须是可序列化的。接收方必须提供同等大小的列表。
+
+dst (int) – 要将 object_list 发送到的目标 rank。目标 rank 基于全局进程组（与 group 参数无关）
+
+group (可选[ProcessGroup]) – (ProcessGroup, 可选): 要在其上工作的进程组。如果为 None，则使用默认进程组。默认为 None。
+
+device (torch.device, 可选) – 如果不为 None，则对象将被序列化并转换为张量，在发送之前移动到该设备。默认为 None。
+
+group_dst (int, 可选) – 组上的目标 rank。必须指定 dst 和 group_dst 中的一个，但不能同时指定两者
+
+use_batch (bool, 可选) – 如果为 True，则使用批量 p2p 操作而不是常规发送操作。这避免了初始化 2 个 rank 的通信器，并使用现有的整个组通信器。有关用法和假设，请参见 batch_isend_irecv。默认为 False。
+
+对于基于 NCCL 的进程组，对象的内部张量表示必须在通信发生之前移动到 GPU 设备。在这种情况下，使用的设备由 torch.cuda.current_device() 给出，用户有责任通过 torch.cuda.set_device() 确保每个 rank 都有一个独立的 GPU。
+
+对象集合通信在性能和可伸缩性方面存在许多严重的限制。详情请参见对象集合通信。
+
+send_object_list() 隐式使用 pickle 模块，众所周知这是不安全的。可能会构造恶意 pickle 数据，在反序列化期间执行任意代码。仅使用您信任的数据调用此函数。
+
+不支持使用 GPU 张量调用 send_object_list() 并且效率低下，因为这会导致 GPU -> CPU 传输（因为张量会被 pickle）。请考虑改用 send()。
+
+同步接收 object_list 中的可序列化对象。
+
+类似于 recv()，但可以接收 Python 对象。
+
+object_list (List[Any]) – 要接收到的对象列表。必须提供与发送列表大小相等的列表。
+
+src (int, 可选) – 要从其接收 object_list 的源 rank。源 rank 基于全局进程组（与 group 参数无关）。如果为 None，将从任何 rank 接收。默认为 None。
+
+group (可选[ProcessGroup]) – (ProcessGroup, 可选): 要在其上工作的进程组。如果为 None，则使用默认进程组。默认为 None。
+
+device (torch.device, 可选) – 如果不为 None，则在此设备上接收。默认为 None。
+
+group_src (int, 可选) – 组上的目标 rank。不能同时指定 src 和 group_src。
+
+use_batch (bool, 可选) – 如果为 True，则使用批量 p2p 操作而不是常规发送操作。这避免了初始化 2 个 rank 的通信器，并使用现有的整个组通信器。有关用法和假设，请参见 batch_isend_irecv。默认为 False。
+
+发送方 rank。如果不属于该组，则为 -1。如果 rank 属于该组，object_list 将包含来自 src rank 发送的对象。
+
+对于基于 NCCL 的进程组，对象的内部张量表示必须在通信发生之前移动到 GPU 设备。在这种情况下，使用的设备由 torch.cuda.current_device() 给出，用户有责任通过 torch.cuda.set_device() 确保每个 rank 都有一个独立的 GPU。
+
+对象集合通信在性能和可伸缩性方面存在许多严重的限制。详情请参见对象集合通信。
+
+recv_object_list() 隐式使用 pickle 模块，众所周知这是不安全的。可能会构造恶意 pickle 数据，在反序列化期间执行任意代码。仅使用您信任的数据调用此函数。
+
+不支持使用 GPU 张量调用 recv_object_list() 并且效率低下，因为这会导致 GPU -> CPU 传输（因为张量会被 pickle）。请考虑改用 recv()。
+
+异步发送或接收一批张量，并返回请求列表。
+
+处理 p2p_op_list 中的每个操作并返回相应的请求。目前支持 NCCL、Gloo 和 UCC 后端。
+
+p2p_op_list (list[torch.distributed.distributed_c10d.P2POp]) – 点对点操作列表（每个运算符的类型为 torch.distributed.P2POp）。列表中 isend/irecv 的顺序很重要，它需要与远程端对应的 isend/irecv 相匹配。
+
+通过调用 op_list 中的相应操作返回的分布式请求对象列表。
+
+list[torch.distributed.distributed_c10d.Work]
+
+请注意，当此 API 与 NCCL PG 后端一起使用时，用户必须使用 torch.cuda.set_device 设置当前的 GPU 设备，否则将导致意外的挂起问题。
+
+此外，如果此 API 是传递给 dist.P2POp 的组中的第一次集合调用，则该组的所有 rank 都必须参与此 API 调用；否则，行为是未定义的。如果此 API 调用不是组中的第一次集合调用，则允许仅涉及该组部分 rank 的批量 P2P 操作。
+
+用于为 batch_isend_irecv 构建点对点操作的类。
+
+这个类构建 P2P 操作的类型、通信缓冲区、
+
+## DistributedDataParallel#
+
+**URL:** https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html
+
+**目录：**
+- DistributedDataParallel#
+
+在模块级别实现基于 `torch.distributed` 的分布式数据并行。
+
+此容器通过同步每个模型副本的梯度来提供数据并行。要进行同步的设备由输入的 `process_group` 指定，默认为整个 world。请注意，DistributedDataParallel 不会在参与的 GPU 之间对输入进行分块或以其他方式进行分片；用户需自行定义如何执行此操作，例如通过使用 `DistributedSampler`。
+
+另见：基础知识和使用 nn.parallel.DistributedDataParallel 代替多进程或 nn.DataParallel。适用于 `torch.nn.DataParallel` 的相同输入约束。
+
+创建此类要求 `torch.distributed` 已经通过调用 `torch.distributed.init_process_group()` 完成初始化。
+
+事实证明，对于单节点多 GPU 数据并行训练，DistributedDataParallel 明显快于 `torch.nn.DataParallel`。
+
+要在具有 N 个 GPU 的主机上使用 DistributedDataParallel，您应该生成 N 个进程，确保每个进程独占地在从 0 到 N-1 的单个 GPU 上工作。这可以通过为每个进程设置 `CUDA_VISIBLE_DEVICES` 或调用以下 GPU API 来完成，
+
+或者调用统一的加速器 API，
+
+其中 i 的范围是 0 到 N-1。在每个进程中，您应参考以下内容来构造此模块：
+
+或者您可以使用最新的 API 进行初始化：
+
+为了在每个节点上生成多个进程，您可以使用 `torch.distributed.launch` 或 `torch.multiprocessing.spawn`。
+
+关于分布式训练的所有功能简介，请参阅 PyTorch 分布式概述。
+
+DistributedDataParallel 可以与 `torch.distributed.optim.ZeroRedundancyOptimizer` 结合使用，以减少每个 rank 的优化器状态内存占用。有关更多详细信息，请参阅 ZeroRedundancyOptimizer 教程。
+
+在使用 GPU 时，nccl 后端是目前最快且强烈推荐的后端。这适用于单节点和多节点分布式训练。
+
+此模块还支持混合精度分布式训练。这意味着您的模型可以拥有不同类型的参数，例如 fp16 和 fp32 的混合类型，对这些混合类型参数的梯度归约也能正常工作。
+
+如果您在一个进程上使用 `torch.save` 对模块进行检查点保存，并在其他进程上使用 `torch.load` 进行恢复，请确保为每个进程正确配置了 `map_location`。如果没有 `map_location`，`torch.load` 会将模块恢复到保存该模块的设备上。
+
+当模型在 M 个节点上以 batch=N 进行训练时，如果损失是在批次中的各个实例间求和（而不是像通常那样求平均），则其梯度将比在单个节点上以 batch=M*N 训练的相同模型小 M 倍（因为不同节点之间的梯度是求平均的）。当您希望获得与本地训练在数学上等价的训练过程时，您应该考虑到这一点。但在大多数情况下，您可以将 DistributedDataParallel 包装的模型、DataParallel 包装的模型和单 GPU 上的普通模型同等对待（例如，在等效批量大小下使用相同的学习率）。
+
+参数永远不会在进程之间广播。该模块会对梯度执行 all-reduce 步骤，并假定它们将在所有进程中以相同的方式被优化器修改。缓冲区（例如 BatchNorm 统计信息）会在每次迭代中从 rank 0 进程的模块广播到系统中的所有其他副本。
+
+如果您将 DistributedDataParallel 与分布式 RPC 框架结合使用，应始终使用 `torch.distributed.autograd.backward()` 来计算梯度，并使用 `torch.distributed.optim.DistributedOptimizer` 来优化参数。
+
+目前，DistributedDataParallel 对使用 `torch.utils.checkpoint()` 的梯度检查点提供有限的支持。如果检查点操作使用 `use_reentrant=False`（推荐），DDP 将按预期工作，没有任何限制。但是，如果检查点操作使用 `use_reentrant=True`（默认值），DDP 只在模型中没有未使用的参数且每层最多只检查点一次时按预期工作（确保不要向 DDP 传递 `find_unused_parameters=True`）。我们目前不支持某层被多次检查点的情况，也不支持检查点模型中存在未使用参数的情况。
+
+为了让非 DDP 模型从 DDP 模型加载状态字典，需要在加载前应用 `consume_prefix_in_state_dict_if_present()` 以去除 DDP 状态字典中的前缀 “module.”。
+
+构造函数、forward 方法以及输出的微分（或此模块输出的函数）是分布式同步点。如果不同的进程可能执行不同的代码，请考虑到这一点。
+
+此模块假定在创建时所有参数都已注册到模型中。以后不应添加或删除任何参数。缓冲区也是如此。
+
+此模块假定每个分布式进程的模型中注册的所有参数顺序都是相同的。模块本身将按照模型注册参数的相反顺序进行梯度 allreduce。换句话说，用户有责任确保每个分布式进程具有完全相同的模型，从而具有完全相同的参数注册顺序。
+
+此模块允许具有非行优先连续步长的参数。例如，您的模型可能包含一些参数的 `torch.memory_format` 为 `torch.contiguous_format`，而另一些参数的格式为 `torch.channels_last`。但是，不同进程中对应的参数必须具有相同的步长。
+
+此模块不能与 `torch.autograd.grad()` 一起使用（即，只有在将梯度累加到参数的 `.grad` 属性中时才起作用）。
+
+如果您计划将此模块与 nccl 后端或（使用 Infiniband 的）gloo 后端一起使用，并与使用多个 worker 的 DataLoader 结合使用，请将多进程启动方法更改为 forkserver（仅限 Python 3）或 spawn。遗憾的是，Gloo（使用 Infiniband）和 NCCL2 不是 fork 安全的，如果不更改此设置，您很可能会遇到死锁。
+
+在用 DistributedDataParallel 包装模型后，绝不应尝试更改模型的参数。因为，当用 DistributedDataParallel 包装模型时，DistributedDataParallel 的构造函数会在构建时为模型本身的所有参数注册额外的梯度归约函数。如果之后更改模型的参数，梯度归约函数将不再匹配正确的参数集。
+
+将 DistributedDataParallel 与分布式 RPC 框架结合使用目前处于实验阶段，可能会发生变化。
+
+module (Module) – 要被并行化的模块
+
+device_ids (list of int 或 torch.device) – CUDA 设备。1) 对于单设备模块，device_ids 只能包含一个设备 id，代表该进程对应的输入模块所在的唯一 CUDA 设备。或者，device_ids 也可以是 None。2) 对于多设备模块和 CPU 模块，device_ids 必须是 None。当在上述两种情况下 device_ids 为 None 时，前向传播的输入数据和实际模块都必须放置在正确的设备上。（默认：None）
+
+CUDA 设备。1) 对于单设备模块，device_ids 只能包含一个设备 id，代表该进程对应的输入模块所在的唯一 CUDA 设备。或者，device_ids 也可以是 None。2) 对于多设备模块和 CPU 模块，device_ids 必须是 None。
+
+当在上述两种情况下 device_ids 为 None 时，前向传播的输入数据和实际模块都必须放置在正确的设备上。（默认：None）
+
+output_device (int 或 torch.device) – 单设备 CUDA 模块的输出设备位置。对于多设备模块和 CPU 模块，它必须是 None，由模块本身决定输出位置。（默认：单设备模块为 device_ids[0]）
+
+broadcast_buffers (bool) – 启用在 forward 函数开始时同步（广播）模块缓冲区的标志。（默认：True）
+
+init_sync (bool) – 是否在初始化期间进行同步，以验证参数形状并广播参数和缓冲区。警告：如果将其设置为 False，用户需自行确保所有 rank 上的权重相同。（默认：True）
+
+process_group – 用于分布式数据 all-reduction 的进程组。如果为 None，将使用由 `torch.distributed.init_process_group()` 创建的默认进程组。（默认：None）
+
+bucket_cap_mb – DistributedDataParallel 会将参数分桶，以便每个桶的梯度归约可以潜在地与反向计算重叠。bucket_cap_mb 控制以 MebiBytes (MiB) 为单位的桶大小。如果为 None，将使用默认大小 25 MiB。（默认：None）
+
+find_unused_parameters (bool) – 从被包装模块的 forward 函数返回值中包含的所有张量遍历 autograd 图。未作为此图的一部分接收梯度的参数将被预先标记为准备好进行归约。此外，在被包装模块的 forward 函数中可能已经使用过但未参与损失计算因而也不会接收梯度的参数，也会被预先标记为准备好进行归约。（默认：False）
+
+check_reduction – 此参数已被弃用。
+
+gradient_as_bucket_view (bool) – 当设置为 True 时，梯度将是指向 allreduce 通信桶不同偏移量的视图。这可以减少峰值内存使用量，节省的内存大小将等于总梯度大小。此外，它避免了梯度和 allreduce 通信桶之间的拷贝开销。当梯度是视图时，不能对梯度调用 `detach_()`。如果遇到此类错误，请参考 `torch/optim/optimizer.py` 中的 `zero_grad()` 函数作为解决方案进行修复。请注意，在第一次迭代后梯度将变成视图，因此应在第一次迭代后检查峰值内存的节省情况。
+
+static_graph (bool) – 当设置为 True 时，DDP 会知道被训练的图是静态的。静态图意味着 1) 已使用和未使用的参数集合在整个训练循环中不会改变；在这种情况下，用户是否设置 `find_unused_parameters = True` 都无关紧要。2) 图的训练方式在整个训练循环中不会改变（意味着不存在取决于迭代的控制流）。当 `static_graph` 设置为 True 时，DDP 将支持过去无法支持的情况：1) 可重入的反向传播。2) 多次激活检查点。3) 模型包含未使用的参数时的激活检查点。4) 存在在 forward 函数之外的模型参数。5) 当存在未使用的参数时可能会提高性能，因为当 `static_graph` 设置为 True 时，DDP 不会在每次迭代中搜索图来检测未使用的参数。要检查是否可以将 `static_graph` 设置为 True，一种方法是在之前的模型训练结束时检查 ddp 日志数据，如果 `ddp_logging_data.get("can_set_static_graph") == True`，通常您也可以将 `static_graph` 设置为 True。示例：>>> model_DDP = torch.nn.parallel.DistributedDataParallel(model) >>> # Training loop >>> ... >>> ddp_logging_data = model_DDP._get_ddp_logging_data() >>> static_graph = ddp_logging_data.get("can_set_static_graph")
+
+当设置为 True 时，DDP 会知道被训练的图是静态的。静态图意味着 1) 已使用和未使用的参数集合在整个训练循环中不会改变；在这种情况下，用户是否设置 `find_unused_parameters = True` 都无关紧要。2) 图的训练方式在整个训练循环中不会改变（意味着不存在取决于迭代的控制流）。当 `static_graph` 设置为 True 时，DDP 将支持过去无法支持的情况：1) 可重入的反向传播。2) 多次激活检查点。3) 模型包含未使用的参数时的激活检查点。4) 存在在 forward 函数之外的模型参数。5) 当存在未使用的参数时可能会提高性能，因为当 `static_graph` 设置为 True 时，DDP 不会在每次迭代中搜索图来检测未使用的参数。要检查是否可以将 `static_graph` 设置为 True，一种方法是在之前的模型训练结束时检查 ddp 日志数据，如果 `ddp_logging_data.get("can_set_static_graph") == True`，通常您也可以将 `static_graph` 设置为 True。
+
+delay_all_reduce_named_params (list of tuple of str 和 torch.nn.Parameter) – 一个命名的参数列表，当 `param_to_hook_all_reduce` 中指定的参数的梯度准备就绪时，这些参数的 all reduce 操作将被延迟。DDP 的其他参数不适用于此参数中指定的命名参数，因为这些命名参数将被 DDP 归约器忽略。
+
+param_to_hook_all_reduce (torch.nn.Parameter) – 用于挂钩 `delay_all_reduce_named_params` 中指定的参数延迟 all reduce 操作的参数。
+
+skip_all_reduce_unused_params – 当设置为 True 时，DDP 将跳过对未使用参数的归约。这要求在整个训练过程中，所有 rank 上的未使用参数保持一致。如果不满足此条件，可能会导致不同步并造成训练挂起。
+
+module (Module) – 要被并行化的模块。
+
+用于在 DDP 中跨进程进行输入不均等训练的上下文管理器。
+
+此上下文管理器将跟踪已加入（joined）的 DDP 进程，并通过插入集合通信操作来“掩盖”前向和反向传播，以与非加入的 DDP 进程创建的操作相匹配。这将确保每个集合通信调用都有一个对应的由已加入的 DDP 进程发出的调用，从而防止在跨进程输入不均等进行训练时可能出现的挂起或错误。或者，如果标志 `throw_on_early_termination` 指定为 True，一旦某个 rank 的输入耗尽，所有训练器都会抛出错误，从而允许根据应用程序逻辑捕获和处理这些错误。
+
+一旦所有 DDP 进程都已加入，上下文管理器会将最后加入的进程对应的模型广播给所有进程，以确保模型在所有进程上都是相同的（这是由 DDP 保证的）。
+
+要使用此功能启用跨进程输入不均等的训练，只需将此上下文管理器包裹在您的训练循环外即可。不需要对模型或数据加载进行进一步的修改。
+
+如果此上下文管理器包裹的模型或训练循环具有额外的分布式集合通信操作，例如模型前向传播中的 `SyncBatchNorm`，则必须启用 `throw_on_early_termination` 标志。这是因为此上下文管理器无法感知非 DDP 的集合通信。当任何一个 rank 的输入耗尽时，此标志将导致所有 rank 抛出异常，从而允许在所有 rank 上捕获这些错误并从中恢复。
+
+divide_by_initial_world_size (bool) – 如果为 True，将使用启动 DDP 训练时的初始 world_size 来除以梯度。如果为 False，将计算有效的 world_size（尚未耗尽其输入的 rank 数量），并在 allreduce 期间除以该数量。设置 `divide_by_initial_world_size=True` 可确保每个输入样本（包括不均等的输入）在对全局梯度的贡献权重上是相等的。这是通过始终使用初始 world_size 来除以梯度来实现的，即使我们遇到不均等的输入也是如此。如果将其设置为 False，我们将使用剩余的节点数来除以梯度。这确保了与较小 world_size 训练的等效性，尽管这意味着不均等的输入会对全局梯度贡献更多。通常，对于训练作业最后几个输入不均等的情况，您会希望将其设置为 True。在输入数量存在巨大差异的极端情况下，将其设置为 False 可能会提供更好的结果。
+
+enable (bool) – 是否启用不均等输入检测。在您知道参与进程之间的输入是均等的情况下，可以传入 `enable=False` 来禁用。默认为 True。
+
+throw_on_early_termination (bool) – 当至少一个 rank 的输入耗尽时，是抛出错误还是继续训练。如果为 True，将在第一个 rank 到达数据末尾时抛出异常。如果为 False，将以较小的有效 world size 继续训练，直到所有 rank 都加入。请注意，如果指定了此标志，则 `divide_by_initial_world_size` 标志将被忽略。默认为 False。
+
+DDP 加入挂钩通过在前向和反向传播中镜像通信来实现输入不均等的训练。
+
+kwargs (dict) – 包含任何用于在运行时修改加入挂钩行为的关键字参数的字典；所有共享相同加入上下文管理器的 Joinable 实例都会接收到相同的 kwargs 值。
+
+如果为 True，则梯度将除以启动 DDP 时使用的初始 world_size。如果为 False，梯度将除以有效的 world_size（即非加入进程的数量），这意味着不均等的输入会对全局梯度贡献更多。通常，如果不均等的程度较小，应将其设置为 True，但在极端情况下，可以将其设置为 False 以获得可能更好的结果。默认为 True。
+
+用于禁用 DDP 进程间梯度同步的上下文管理器。
+
+在此上下文中，梯度将累积在模块变量上，随后将在退出上下文后的第一次前向-反向传播中进行同步。
+
+前向传播应包含在上下文管理器内部，否则梯度仍将被同步。
+
+为用户自定义的跨多个 worker 的 DDP 梯度聚合注册通信挂钩。
+
+这个挂钩对于研究人员尝试新想法非常有用。例如，此挂钩可用于实现诸如 GossipGrad 和梯度压缩等几种算法，这些算法在运行分布式数据并行训练时涉及不同的参数同步通信策略。
+
+state (object) – 传递给挂钩以在训练过程中维护任何状态信息。示例包括梯度压缩中的误差反馈、GossipGrad 中下一个要通信的节点等。它由每个 worker 在本地存储，并由该 worker 上的所有梯度张量共享。
+
+传递给挂钩以在训练过程中维护任何状态信息。示例包括梯度压缩中的误差反馈、GossipGrad 中下一个要通信的节点等。
+
+它由每个 worker 在本地存储，并由该 worker 上的所有梯度张量共享。
+
+hook (Callable) – 具有以下签名的可调用对象：`hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]`：此函数在桶准备就绪时被调用。挂钩可以执行所需的任何处理，并返回一个指示任何异步工作（如 allreduce）完成的 Future。如果挂钩不执行任何通信，它仍必须返回一个已完成的 Future。Future 应包含梯度桶张量的新值。一旦桶准备就绪，c10d 归约器将调用此挂钩，并使用 Future 返回的张量将梯度复制到各个参数。请注意，future 的返回类型必须是单个张量。我们还提供了一个名为 `get_future` 的 API，用于检索与 `c10d.ProcessGroup.Work` 完成相关联的 Future。`get_future` 目前支持 NCCL，并且也支持 GLOO 和 MPI 上的大多数操作，但不支持点对点操作（send/recv）。
+
+具有以下签名的可调用对象：`hook(state: object, bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]`：
+
+此函数在桶准备就绪时被调用。挂钩可以执行所需的任何处理，并返回一个指示任何异步工作（如 allreduce）完成的 Future。如果挂钩不执行任何通信，它仍必须返回一个已完成的 Future。Future 应包含梯度桶张量的新值。一旦桶准备就绪，c10d 归约器将调用此挂钩，并使用 Future 返回的张量将梯度复制到各个参数。请注意，future 的返回类型必须是单个张量。
+
+我们还提供了一个名为 `get_future` 的 API，用于检索与 `c10d.ProcessGroup.Work` 完成相关联的 Future。`get_future` 目前支持 NCCL，并且也支持 GLOO 和 MPI 上的大多数操作，但不支持点对点操作（send/recv）。
+
+梯度桶的张量不会预先除以 world_size。在进行诸如 allreduce 之类的操作时，用户有责任自行除以 world_size。
+
+DDP 通信挂钩只能注册一次，并且应该在调用 backward 之前进行注册。
+
+挂钩返回的 Future 对象应包含一个与梯度桶内张量形状相同的单个张量。
+
+`get_future` API 支持 NCCL，并部分支持 GLOO 和 MPI 后端（不支持诸如 send/recv 之类的点对点操作），并且将返回一个 `torch.futures.Future`。
+
+下面是一个返回相同张量的无操作挂钩示例。
+
+下面是一个并行 SGD 算法的示例，其中梯度在 allreduce 之前被编码，然后在 allreduce 之后被解码。
+
+---
+
+## DDP Communication Hooks#
+
+**URL:** https://pytorch.org/docs/stable/ddp_comm_hooks.html
+
+**目录：**
+- DDP Communication Hooks#
+- How to Use a Communication Hook?#
+- What Does a Communication Hook Operate On?#
+- 默认通信挂钩#
+- PowerSGD Communication Hook#
+  - PowerSGD State#
+  - PowerSGD Hooks#
+- Debugging Communication Hooks#
+- Checkpointing of Communication Hooks#
+- Acknowledgements#
+
+Created On: Jun 06, 2025 | Last Updated On: Jun 06, 2025
+
+DDP 通信挂钩是一个通用接口，用于通过覆盖 DistributedDataParallel 中原生的 allreduce 来控制如何跨 worker 通信梯度。系统提供了几个内置的通信挂钩，用户可以轻松应用这些挂钩中的任何一个来优化通信。此外，该挂钩接口还可以为更高级的用例支持用户自定义的通信策略。
+
+要使用通信挂钩，用户只需在训练循环之前让 DDP 模型注册该挂钩即可，如下所示。
+
+torch.nn.parallel.DistributedDataParallel.register_comm_hook()
+
+通信挂钩提供了一种灵活的方式来进行梯度的 allreduce。因此，它主要在 allreduce 之前对每个副本上的梯度进行操作，这些梯度被分桶以增加通信和计算之间的重叠。特别地，`torch.distributed.GradBucket` 表示一桶需要进行 allreduce 的梯度张量。
+
+此类主要将扁平化的梯度张量（由 `buffer()` 返回）传递给 DDP 通信挂钩。该张量可以进一步分解为该桶内每个参数的独立张量列表（由 `get_per_parameter_tensors()` 返回），以应用逐层操作。
+
+由于桶在第一次迭代后会被重建，因此不应在训练开始时依赖这些索引。
+
+存储少数几层连续梯度的桶的索引。所有梯度都会被分桶。
+
+一个扁平化的 1D `torch.Tensor` 缓冲区，可以将其进一步分解为该桶内每个参数的独立张量列表。
+
+一个 `torch.Tensor` 列表。列表中的每个张量对应一个梯度。
+
+此桶是否是迭代中最后一个进行 allreduce 的桶。这也意味着此桶对应于前向传播中的最前面几层。
+
+使用输入张量缓冲区替换桶中的张量。
+
+一个 `torch.Tensor` 列表。列表中的每个张量对应一个模型参数。
+
+默认的通信挂钩是简单的无状态挂钩，因此 `register_comm_hook` 中的输入状态是进程组或 None。输入桶是一个 `torch.distributed.GradBucket` 对象。
+
+使用 GradBucket 张量调用 allreduce。
+
+一旦梯度张量在所有 worker 之间聚合完成，它的回调函数就会取平均值并返回结果。
+
+如果用户注册了此 DDP 通信挂钩，DDP 的结果预期将与未注册挂钩的情况相同。因此，这不会改变 DDP 的行为，用户可以将其作为参考，或修改此挂钩以记录有用的信息或在不影响 DDP 行为的情况下用于任何其他目的。
+
+通过将 GradBucket 强制转换为 `torch.float16` 并除以进程组大小来进行压缩。
+
+此 DDP 通信挂钩实现了一种简单的梯度压缩方法，将 GradBucket 张量转换为半精度浮点格式 (`torch.float16`)，然后将其除以进程组大小。它对这些 float16 梯度张量进行 allreduce。一旦压缩的梯度张量被 allreduce，链式回调的解压操作会将其转换回输入数据类型（例如 float32）。
+
+警告：此 API 是实验性的，并且要求 NCCL 版本高于 2.9.6。
+
+此 DDP 通信挂钩实现了一种简单的梯度压缩方法，将 GradBucket 张量转换为半精度 Brain 浮点格式 (`torch.bfloat16`)，然后将其除以进程组大小。它对这些 bfloat16 梯度张量进行 allreduce。一旦压缩的梯度张量被 allreduce，链式回调的解压操作会将其转换回输入数据类型（例如 float32）。
+
+此外，还提供了一个通信挂钩包装器，以支持将 `fp16_compress_hook()` 或 `bf16_compress_hook()` 作为包装器，这可以与其他通信挂钩结合使用。
+
+将输入张量转换为 `torch.float16`，将挂钩的结果转换回输入数据类型。
+
+此包装器将给定 DDP 通信挂钩的输入梯度张量转换为半精度浮点格式 (`torch.float16`)，并将给定挂钩的结果张量转换回输入数据类型，例如 float32。因此，`fp16_compress_hook` 等同于 `fp16_compress_wrapper(allreduce_hook)`。
+
+Callable[[Any, GradBucket], Future[Tensor]]
+
+警告：此 API 是实验性的，并且要求 NCCL 版本高于 2.9.6。
+
+此包装器将给定 DDP 通信挂钩的输入梯度张量转换为半精度 Brain 浮点格式 (`torch.bfloat16`)，并将给定挂钩的结果张量转换回输入数据类型，例如 float32。
+
+因此，`bf16_compress_hook` 等同于 `bf16_compress_wrapper(allreduce_hook)`。
+
+Callable[[Any, GradBucket], Future[Tensor]]
+
+PowerSGD（Vogels 等人，NeurIPS 2019）是一种梯度压缩算法，可以提供非常高的压缩率并加速受限于带宽的分布式训练。此算法需要同时维护一些超参数和内部状态。因此，PowerSGD 通信挂钩是一个有状态挂钩，用户需要提供如下定义的状态对象。
+
+存储训练期间所有梯度的算法超参数和内部状态。
+
+特别是，`matrix_approximation_rank` 和 `start_powerSGD_iter` 是需要用户调整的主要超参数。为了性能，我们建议保持二进制超参数 `use_error_feedback` 和 `warm_start` 为开启状态。
+
+`matrix_approximation_rank` 控制压缩低秩张量的大小，这决定了压缩率。秩越低，压缩越强。
+
+1.1. 如果 `matrix_approximation_rank` 太低，完整的模型质量将需要更多的训练步骤才能达到，或者永远无法达到，从而导致精度损失。
+
+1.2. `matrix_approximation_rank` 的增加会大幅增加压缩的计算成本，并且精度可能不会超过某个 `matrix_approximation_rank` 阈值而进一步提高。
+
+要调整 `matrix_approximation_rank`，我们建议从 1 开始并以 2 的倍数增加（如指数网格搜索，1、2、4...），直到达到令人满意的精度。通常只使用较小的值 1-4。对于某些 NLP 任务（如原论文的附录 D 所示），此值已增加到 32。
+
+`start_powerSGD_iter` 将 PowerSGD 压缩推迟到步骤 `start_powerSGD_iter`，而在步骤 `start_powerSGD_iter` 之前运行原生的 allreduce。这种原生 allreduce + PowerSGD 的混合方案可以有效提高精度，即使使用相对较小的 `matrix_approximation_rank` 也是如此。这是因为训练阶段的开头通常对不准确的梯度非常敏感，过早压缩梯度可能会使训练迅速进入次优轨迹，从而对精度造成无法恢复的影响。
+
+要调整 `start_powerSGD_iter`，我们建议从总训练步数的 10% 开始，然后增加直到达到令人满意的精度。如果训练中包含预热阶段，`start_powerSGD_iter` 通常不应小于预热步数。
+
+`min_compression_rate` 是压缩层时所需的最小压缩率。由于压缩会带来计算开销，只有在能够充分节省带宽的情况下，才值得压缩张量，即满足 `(num_rows + num_cols) * matrix_approximation_rank * min_compression_rate < num_rows * num_cols`。如果无法满足指定的压缩率阈值，则将直接对该张量进行 allreduce 而不进行压缩。
+
+一旦 PowerSGD 压缩开始，将每隔 `compression_stats_logging_frequency` 次迭代记录一次压缩统计数据。
+
+`orthogonalization_epsilon` 可以是一个非常小的值（例如 1e-8），在正交化步骤中将其添加到每个归一化矩阵列中，以防止在任何列全为 0 时出现除以零错误。如果这已经被防止了（例如通过批归一化），为了准确性，建议将 epsilon 设置为 0。
+
+`batch_tensors_with_same_shape` 控制是否在批处理操作中压缩和解压具有相同形状的张量，以实现更高的并行度。请注意，您还应该增加桶大小（即 DDP 构造函数中的 `bucket_cap_mb` 参数），以使更多具有相同形状的张量出现在同一个桶中，然而这可能会减少计算和通信之间的重叠，并由于堆叠相同形状的张量而增加内存占用。如果压缩/解压计算成为瓶颈，请设置为 True。
+
+如果启用了误差反馈或预热，DDP 中允许的 `start_powerSGD_iter` 的最小值为 2。这是因为 DDP 中还有另一个内部优化会在第 1 次迭代时重建桶，这可能会与在重建过程之前记忆的任何张量发生冲突。
+
+PowerSGD 通常需要与模型梯度大小相同的额外内存来启用误差反馈，这可以补偿有偏的压缩通信并提高精度。
+
+PowerSGD 挂钩可能与 Apex 自动混合精度包冲突。请改用 PyTorch 原生的自动混合精度包。
+
+实现 PowerSGD 算法。
+
+此 DDP 通信挂钩实现了论文中描述的 PowerSGD 梯度压缩算法。一旦梯度张量在所有 worker 之间聚合完成，此挂钩将按如下方式应用压缩：
+
+将输入的扁平化 1D 梯度张量视为每个参数的独立张量列表，并将所有张量分为两组：
+
+1.1 在 allreduce 之前应该被压缩的张量，因为压缩可以节省足够的带宽。
+
+1.2 其余张量将被直接进行 allreduce 而不进行压缩，包括所有的向量张量（例如偏置）。
+
+处理未压缩的张量：
+
+2.1. 为这些未压缩的张量分配连续的内存，并将所有未压缩的张量作为一个批次进行 allreduce，不进行压缩；
+
+2.2. 将单个未压缩的张量从连续内存复制回输入张量。
+
+处理应被 PowerSGD 压缩的张量：
+
+3.1. 对于每个张量 M，创建两个低秩张量 P 和 Q 用于分解 M，使得 M = PQ^T，其中 Q 从标准正态分布初始化并进行正交化；
+
+3.2. 计算 Ps 中的每个 P，等于 MQ；
+
+3.3. 将 Ps 作为一个批次进行 allreduce；
+
+3.4. 对 Ps 中的每个 P 进行正交化；
+
+3.5. 计算 Qs 中的每个 Q，近似等于 M^TP；
+
+3.6. 将 Qs 作为一个批次进行 allreduce；
+
+3.7. 计算所有压缩张量中的每个 M，近似等于 PQ^T。
+
+请注意，此通信挂钩在前 `state.start_powerSGD_iter` 次迭代中强制使用原生 allreduce。这不仅让用户能更好地控制速度和精度之间的权衡，也有助于为未来的通信挂钩开发者抽象掉 DDP 内部优化的某些复杂性。
+
+state (PowerSGDState) – 用于配置压缩率并支持误差反馈、热启动等的状态信息。要调整压缩配置，主要需要调整 `matrix_approximation_rank`、`start_powerSGD_iter` 和 `min_compression_rate`。
+
+bucket (dist.GradBucket) – 存储扁平化 1D 梯度张量（打包了多个按变量划分的张量）的桶。请注意，由于 DDP 通信挂钩仅支持单进程单设备模式，因此该桶中仅存储了刚好一个张量。
+
+通信的 Future 处理程序，它会原地更新梯度。
+
+实现简化的 PowerSGD 算法。
+
+此 DDP 通信挂钩实现了论文中描述的简化版 PowerSGD 梯度压缩算法。此变体不逐层压缩梯度，而是压缩打包了所有梯度的扁平化输入张量。因此，它比 `powerSGD_hook()` 更快，但通常会导致精度低得多，除非 `matrix_approximation_rank` 为 1。
+
+在这里增加 `matrix_approximation_rank` 不一定会提高精度，因为在不进行列/行对齐的情况下打包每个参数的张量可能会破坏低秩结构。因此，用户应始终优先考虑 `powerSGD_hook()`，只有当 `matrix_approximation_rank` 为 1 就能达到满意的精度时，才考虑此变体。
+
+一旦梯度张量在所有 worker 之间聚合完成，此挂钩将按如下方式应用压缩：
+
+将输入的扁平化 1D 梯度张量视为带 0 填充的方形张量 M；
+
+创建两个低秩张量 P 和 Q 用于分解 M，使得 M = PQ^T，其中 Q 从标准正态分布初始化并进行正交化；
+
+计算 P，等于 MQ；
+
+计算 Q，近似等于 M^TP；
+
+计算 M，近似等于 PQ^T。
+
+将输入张量截断为原始长度。
+
+请注意，此通信挂钩在前 `state.start_powerSGD_iter` 次迭代中强制使用原生 allreduce。这不仅让用户能更好地控制速度和精度之间的权衡，也有助于为未来的通信挂钩开发者抽象掉 DDP 内部优化的某些复杂性。
+
+state (PowerSGDState) – 用于配置压缩率并支持误差反馈、热启动等的状态信息。要调整压缩配置，主要需要调整 `matrix_approximation_rank` 和 `start_powerSGD_iter`。
+
+bucket (dist.GradBucket) – 存储扁平化 1D 梯度张量（打包了多个按变量划分的张量）的桶。请注意，由于 DDP 通信挂钩仅支持单进程单设备模式，因此该桶中仅存储了刚好一个张量。
+
+通信的 Future 处理程序，它会原地更新梯度。
+
+顾名思义，调试通信挂钩仅用于调试和性能优化的目的。
+
+调试通信挂钩不一定会输出正确的结果。
+
+返回一个包装了输入的 future，因此它是一个不会产生任何通信开销的无操作。
+
+此挂钩应仅用于 allreduce 优化的裕度分析，而不是用于常规的梯度同步。例如，如果注册此挂钩后只能观察到不到 10% 的训练时间加速，通常意味着 allreduce 对于此情况不是性能瓶颈。如果无法轻易获取 GPU 追踪或追踪分析由于 allreduce 和计算之间的重叠或各 rank 之间的不同步等某些因素而变得复杂时，这种检测会特别有用。
+
+有状态的通信挂钩可以作为模型检查点的一部分进行保存，以启用训练器重启。要使挂钩可序列化，应定义 `__setstate__` 和 `__getstate__`。
+
+`__getstate__` 应从返回的字典中排除不可序列化的属性。
+
+`__setstate__` 应正确初始化所提供状态中排除的不可序列化属性。
+
+PowerSGDState 已经实现了 `__setstate__` 和 `__getstate__`，可以用作参考。
+
+返回一个将被序列化并保存的 `Dict[str, Any]`。
+
+`process_group` 是不可序列化的，已从返回的状态中排除。
+
+接收提供的状态并设置到此 PowerSGDState 实例。
+
+`process_group` 被设置为默认值。
+
+这里有一个保存和重新加载 PowerSGD 状态及挂钩的简单、端到端示例。
+
+非常感谢 PowerSGD 论文的作者 Thijs Vogels 对 PowerSGD 通信挂钩的代码审查，以及对比实验，这表明 PowerSGD 通信挂钩的性能与原论文中的实现不相上下。
+
+## 分布式检查点 - torch.distributed.checkpoint#
+
+**URL:** https://pytorch.org/docs/stable/distributed.checkpoint.html
+
+**目录：**
+- 分布式检查点 - torch.distributed.checkpoint#
+- 附加资源:#
+
+创建于：2022年11月16日 | 最后更新于：2025年9月4日
+
+分布式检查点 (DCP) 支持从多个 rank 并行加载和保存模型。它处理加载时的重新分片，这使得在一个集群拓扑中保存并在另一个集群拓扑中加载成为可能。
+
+DCP 在几个重要方面与 `torch.save` 和 `torch.load` 不同：
+
+每次检查点它会生成多个文件，每个 rank 至少生成一个。
+
+它是原地操作的，这意味着模型应该首先分配其数据，然后 DCP 使用该存储空间。
+
+加载和保存检查点的入口点如下：
+
+分布式检查点 (DCP) 入门
+
+使用分布式检查点 (DCP) 进行异步保存
+
+TorchTitan 检查点文档
+
+TorchTitan DCP 实现
+
+用于异步检查点类型的枚举。
+
+此类包含用于暂存和上传完成的 Future 对象。它由 `async_save()` 返回。`staging_completion` 是一个 Future，指示 `state_dict` 的本地副本何时完成。`upload_completion` 是一个 Future，指示检查点何时完成保存。
+
+以 SPMD（单程序多数据）风格保存分布式模型。
+
+此函数与 `torch.save()` 不同，因为它处理 `ShardedTensor` 和 `DTensor`，方式是让每个 rank 仅保存其本地分片。
+
+对于每个有状态对象（同时具有 `state_dict` 和 `load_state_dict`），保存时将在序列化之前调用 `state_dict`。
+
+对于保存的 `state_dict`，不保证跨 PyTorch 版本的向后兼容性。
+
+如果使用 `process_group` 参数，请确保只有其 rank 调用 `save_state_dict`，并且 `state_dict` 中的所有数据都属于该进程组。
+
+当为 FSDP 的 `ShardingStrategy.HYBRID_SHARD` 保存检查点时，`shard_group` 中应只有一个在调用 `save_state_dict`，并且需要传入相应的进程组。
+
+本地进程中的 `state_dict`。
+
+state_dict (Dict[str, Any]) – 要保存的 state_dict。
+
+checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是一个键。(默认: None)
+
+storage_writer (可选[StorageWriter]) – 用于执行写入操作的 `StorageWriter` 实例。如果未指定，DCP 将根据 `checkpoint_id` 自动推断写入器。如果 `checkpoint_id` 也为 None，则会引发异常。(默认: None)
+
+planner (可选[SavePlanner]) – `SavePlanner` 的实例。如果未指定，将使用默认规划器。(默认: None)
+
+process_group (可选[ProcessGroup]) – 用于跨 rank 同步的进程组。(默认: None)
+
+no_dist (bool) – 如果为 True，此函数将假定意图是在单个 rank/进程上加载检查点。(默认: False)
+
+use_collectives (bool) – 如果为 False，此函数将假定意图是在不使用跨 rank 同步的情况下保存检查点。(默认: True) 此配置是实验性的，应谨慎使用。它会更改保存的检查点的格式，并且可能无法向后兼容。
+
+保存的检查点的元数据对象。
+
+`save_state_dict` 使用集合通信来协调跨 rank 的写入。对于基于 NCCL 的进程组，对象的内部张量表示必须在通信发生之前移动到 GPU 设备。在这种情况下，所使用的设备由 `torch.cuda.current_device()` 给出，用户有责任通过 `torch.cuda.set_device()` 确保已设置此项，以便每个 rank 都有一个独立的 GPU。
+
+`save` 的异步版本。此代码首先将 `state_dict` 暂存到暂存存储（默认为 CPU 内存），然后在一个单独的线程中调用保存。
+
+此功能是实验性的，可能会发生变化。在保存最后一个检查点后必须调用 CLOSE。
+
+state_dict (Dict[str, Any]) – 要保存的 state_dict。
+
+checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是一个键。(默认: None)
+
+storage_writer (可选[StorageWriter]) – 用于执行“暂存”和“保存”的 `StorageWriter` 实例。如果未指定，DCP 将根据 `checkpoint_id` 自动推断写入器。如果 `checkpoint_id` 也为 None，则会引发异常。(默认: None)
+
+planner (可选[SavePlanner]) – `SavePlanner` 的实例。如果未指定，将使用默认规划器。(默认: None)
+
+process_group (可选[ProcessGroup]) – 用于跨 rank 同步的进程组。(默认: None)
+
+async_checkpointer_type (AsyncCheckpointerType) – 是在单独的线程还是进程中执行检查点操作 (默认: AsyncCheckpointerType.THREAD)
+
+async_stager (AsyncStager) – 提供暂存实现。如果 `storage_writer` 实现了 `AsyncStager` 并且提供了 `async_stager`，则将使用 `async_stager` 进行暂存
+
+no_dist (bool) – 如果为 True，此函数将假定意图是在单个 rank/进程上保存检查点。(默认: False)
+
+use_collectives (bool) – 如果为 False，在没有 rank 协调的情况下保存检查点。(默认: True) 此配置是实验性的，应谨慎使用。它会更改保存的检查点的格式，并且可能无法向后兼容。
+
+包含来自保存操作生成的 Metadata 对象的 Future。
+
+此方法已弃用。请切换到 ‘save’。
+
+以 SPMD 风格将检查点加载到分布式 state dict 中。
+
+每个 rank 在提供给此 API 的 `state_dict` 中必须具有相同的键。不匹配的键可能会导致挂起或错误。如果不确定，可以使用 `utils._assert_same_keys` API 进行检查（但可能会产生通信成本）。
+
+每个 rank 将尝试读取最少量的必要数据，以满足所请求的 `state_dict`。当加载 `ShardedTensor` 或 `DTensor` 实例时，每个 rank 仅读取其本地分片的数据。
+
+对于每个有状态对象（同时具有 `state_dict` 和 `load_state_dict`），加载将首先在尝试反序列化之前调用 `state_dict`，并在反序列化完成后调用 `load_state_dict`。对于每个非有状态对象，加载将对该对象进行反序列化，然后用反序列化后的对象替换 `state_dict` 中的该对象。
+
+在调用此函数之前，`state_dict` 中的所有张量都必须分配到其目标设备上。
+
+所有非张量数据都使用 `torch.load()` 加载，并在 `state_dict` 中进行原位修改。
+
+用户必须在根模块上调用 `load_state_dict`，以确保加载后处理和非张量数据的正确传播。
+
+state_dict (Dict[str, Any]) – 要将检查点加载到的 state_dict。
+
+checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是一个键。(默认: None)
+
+storage_reader (可选[StorageReader]) – 用于执行读取操作的 `StorageReader` 实例。如果未指定，DCP 将根据 `checkpoint_id` 自动推断读取器。如果 `checkpoint_id` 也为 None，则会引发异常。(默认: None)
+
+planner (可选[LoadPlanner]) – `LoadPlanner` 的实例。如果未指定，将使用默认规划器。(默认: None)
+
+process_group (可选[ProcessGroup]) – 用于跨 rank 同步的进程组。(默认: None)
+
+no_dist (bool) – 如果为 True，此函数将假定意图是在不使用跨 rank 同步的情况下加载检查点。(默认: False)
+
+`load_state_dict` 使用集合通信来协调跨 rank 的读取。对于基于 NCCL 的进程组，对象的内部张量表示必须在通信发生之前移动到 GPU 设备。在这种情况下，所使用的设备由 `torch.cuda.current_device()` 给出，用户有责任通过 `torch.cuda.set_device()` 确保已设置此项，以便每个 rank 都有一个独立的 GPU。
+
+此方法已弃用。请切换到 ‘load’。
+
+以下模块对于进一步自定义用于异步检查点的暂存机制 也很有用：
+
+此协议旨在为 `dcp.async_save` 提供自定义和可扩展性，允许用户自定义在并行执行常规 `dcp.save` 路径之前如何暂存数据。预期的操作顺序（具体在 `torch.distributed.state_dict_saver.async_save` 中定义）如下：
+
+此调用给 `AsyncStager` 提供“暂存” `state_dict` 的机会。在此上下文中，暂存的期望和目的是创建 state dict 的“训练安全”表示，这意味着在暂存完成后对模块数据的任何更新都不应反映在此方法返回的 state dict 中。例如，在默认情况下，会在 CPU RAM 上创建整个 state dict 的副本并在此处返回，允许用户继续训练而不会冒正在序列化的数据被更改的风险。
+
+用于序列化 `state_dict` 并将其写入存储。
+
+序列化线程启动之前以及从 `dcp.async_save` 返回之前。如果将其设置为 False，则假定用户已经定义了一个自定义同步点，目的是为了进一步优化训练循环中的保存延迟（例如，将暂存与前向/后向传递重叠），并且用户有责任在适当的时间调用 `AsyncStager.synchronize_staging`。
+
+清理暂存器使用的所有资源。
+
+执行暂存后是否同步。
+
+返回 `state_dict` 的“暂存”副本。对暂存副本的期望是，它不受在暂存调用完成后发生的任何更新的影响。
+
+Union[Future[dict[str, Union[~StatefulT, Any]]], dict[str, Union[~StatefulT, Any]]]
+
+如果暂存以某种方式是异步的，则应调用此方法以确保暂存完成，并且可以安全地开始修改原始的 `state_dict`。
+
+`DefaultStager` 提供了一个功能齐全的暂存实现，它结合了多种优化技术以高效地准备检查点。
+
+暂存过程的工作原理如下： 1. 提交状态字典以进行暂存（同步或异步） 2. 将张量从 GPU 复制到优化的 CPU 存储中 3. 如果使用非阻塞复制，则同步 CUDA 操作 4. 通过 Future 返回或提供已暂存的状态字典
+
+# 同步暂存 stager = DefaultStager(StagingOptions(use_async_staging=False)) staged_dict = stager.stage(state_dict) stager.close()
+
+# 异步暂存 stager = DefaultStager(StagingOptions(use_async_staging=True)) future = stager.stage(state_dict) # …… 执行其他工作 …… staged_dict = future.result() stager.close()
+
+# 上下文管理器模式（推荐） stager = DefaultStager(config) with stager: result = stager.stage(state_dict)
+
+当模型计算可以与暂存操作重叠时，异步暂存可提供最佳性能。
+
+固定内存可提高 CPU-GPU 传输速度，但会使用更多的内存。
+
+共享内存允许与检查点进程进行高效的 IPC（进程间通信）。
+
+非阻塞复制减少了内存传输期间的 GPU 空闲时间。
+
+`DefaultStager` 不是线程安全的。每个线程应使用自己的实例，或者应提供外部同步。
+
+清理 `DefaultStager` 使用的所有资源。关闭用于异步暂存操作的 `ThreadPoolExecutor`，并清理底层 `StateDictStager` 的缓存存储。当不再需要暂存器时应调用此方法以防止资源泄漏，特别是在长时间运行的应用程序中。调用 `close()` 后，不应将其用于进一步的暂存操作。
+
+stager = DefaultStager(StagingOptions(use_async_staging=True)) future = stager.stage(state_dict) result = future.result() stager.close() # 清理所有资源
+
+此函数负责暂存 `state_dict`。有关暂存的更多详细信息，请参见类文档字符串。如果 `use_async_staging` 为 True，它将返回一个在暂存完成时会被履行的 Future 对象。如果 `use_async_staging` 为 False，它将返回完全暂存的 `state_dict`。
+
+state_dict (STATE_DICT_TYPE) – 要暂存的 state_dict。
+
+Union[dict[str, Union[~StatefulT, Any]], Future[dict[str, Union[~StatefulT, Any]]]]
+
+当 `use_async_staging` 为 True 时，此方法将等待直到暂存完成。如果 `use_async_staging` 为 False，则此方法为无操作。
+
+配置检查点暂存行为的选项。
+
+use_pinned_memory (bool) – 启用固定内存分配以实现更快的 CPU-GPU 传输。需要 CUDA 可用。 默认: True
+
+use_shared_memory (bool) – 为多进程场景启用共享内存。当多个进程需要访问相同的暂存数据时非常有用。 默认: True
+
+use_async_staging (bool) – 使用后台线程池启用异步暂存。允许将计算与暂存操作重叠。需要 CUDA。 默认: True
+
+use_non_blocking_copy (bool) – 使用带有流同步的非阻塞设备内存复制。通过允许 CPU 工作在 GPU 传输期间继续进行来提高性能。 默认: True
+
+如果 CUDA 不可用，依赖于 CUDA 的功能将引发异常。
+
+`AsyncStager` 的一种实现，它在 CPU RAM 上暂存 `state_dict` 并阻塞直到复制完成。此实现还提供了一个选项，以使用固定内存优化暂存延迟。
+
+注意：在这种情况下 `synchronize_staging` 是一个无操作的方法。
+
+返回 CPU 上的 `state_dict` 副本。
+
+dict[str, Union[~StatefulT, Any]]
+
+无操作函数，因为暂存是阻塞的。
+
+除了上述入口点之外，如下所述的有状态对象在保存/加载期间提供了额外的自定义功能。
+
+用于可以被检查点和恢复的对象的有状态协议。
+
+从提供的 `state_dict` 恢复对象的状态。
+
+state_dict (dict[str, Any]) – 要从中恢复的 state dict。
+
+对象应将其 `state_dict` 表示形式作为字典返回。此函数的输出将被检查点保存，并稍后在 `load_state_dict()` 中恢复。
+
+由于恢复检查点的原位特性，此函数也会在 `torch.distributed.checkpoint.load` 期间被调用。
+
+对象状态字典。
+
+此示例演示如何使用 Pytorch 分布式检查点保存 FSDP 模型。
+
+以下类型定义了检查点期间使用的 IO 接口：
+
+`load_state_dict` 用来从存储中读取的接口。
+
+一个 `StorageReader` 实例在分布式检查点中同时充当协调者和跟随者。作为初始化的一部分，每个实例都会被告知其角色。
+
+子类应预期 `load_state_dict` 会进行以下调用序列：
+
+（所有 rank）如果用户传递了有效的 `checkpoint_id`，则设置 `checkpoint_id`。
+
+（所有 rank）`read_metadata()`
+
+（所有 rank）`set_up_storage_reader()`
+
+（所有 rank）`prepare_local_plan()`
+
+（协调者）`prepare_global_plan()`
+
+（所有 rank）`read_data()`
+
+执行存储加载的集中规划。
+
+此方法仅在协调者实例上调用。
+
+虽然此方法可以产生一个完全不同的计划，但首选方法是将存储特定数据存储在 `LoadPlan::storage_data` 中。
+
+plans (list[torch.distributed.checkpoint.planner.LoadPlan]) – 一个 `LoadPlan` 实例列表，每个 rank 一个。
+
+存储全局规划后转换后的 `LoadPlan` 列表
+
+list[torch.distributed.checkpoint.planner.LoadPlan]
+
+执行特定于存储的本地规划。
+
+虽然此方法可以产生一个完全不同的计划，但推荐的方法是将存储特定数据存储在 `LoadPlan::storage_data` 中。
+
+plan (LoadPlan) – 使用中的 `LoadPlan` 的本地计划。
+
+存储本地规划后转换后的 `LoadPlan`
+
+使用规划器解析数据，并从计划中读取所有项目。
+
+子类应调用 `LoadPlanner::load_bytes` 将 `BytesIO` 对象反序列化到正确的位置。
+
+子类应调用 `LoadPlanner::resolve_tensor` 以获取应该将数据加载到的张量的访问权限。
+
+`StorageLayer` 有责任正确调度任何所需的跨设备复制。
+
+plan (LoadPlan) – 要执行的本地计划
+
+planner (LoadPlanner) – 用于解析项目的规划器对象。
+
+一个在所有读取完成时结束的 Future。
+
+读取检查点元数据。
+
+与正在加载的检查点关联的元数据对象。
+
+调用表示即将进行一次全新的检查点读取。如果用户为此检查点读取设置了 `checkpoint_id`，则可能会存在该 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹/文件的路径，也可以是键值存储的键。
+
+checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹或文件的路径。如果存储更像是一个键值存储，它也可以是一个键。(默认: None)
+
+初始化此实例。
+
+metadata (Metadata) – 要使用的元数据架构。
+
+is_coordinator (bool) – 此实例是否负责协调检查点。
+
+检查给定的 `checkpoint_id` 是否受存储支持。这允许我们启用自动存储选择。
+
+`save_state_dict` 用来写入存储的接口。
+
+一个 `StorageWriter` 实例在分布式检查点中同时充当协调者和跟随者。作为初始化的一部分，每个实例都会被告知其角色。
+
+子类应预期会进行以下调用序列。
+
+（所有 rank）如果用户传递了有效的 `checkpoint_id`，则设置 `checkpoint_id`。
+
+（所有 rank）`set_up_storage_writer()`
+
+（所有 rank）`prepare_local_plan()`
+
+（协调者）`prepare_global_plan()`
+
+（所有 rank）`write_data()`
+
+（协调者）`finish()`
+
+写入元数据并标记当前检查点为成功。
+
+用于序列化元数据的实际格式/架构是一个实现细节。唯一的要求是它可以恢复为相同的对象图。
+
+metadata (Metadata) – 新检查点的元数据
+
+results (list[list[torch.distributed.checkpoint.storage.WriteResult]]) – 来自所有 rank 的 `WriteResult` 列表。
+
+执行存储的集中规划。
+
+此方法仅在协调者实例上调用。
+
+虽然此方法可以产生一个完全不同的计划，但首选方法是将存储特定数据存储在 `SavePlan::storage_data` 中。
+
+plans (list[torch.distributed.checkpoint.planner.SavePlan]) – 一个 `SavePlan` 实例列表，每个 rank 一个。
+
+存储全局规划后转换后的 `SavePlan` 列表
+
+list[torch.distributed.checkpoint.planner.SavePlan]
+
+执行特定于存储的本地规划。
+
+虽然此方法可以产生一个完全不同的计划，但推荐的方法是将存储特定数据存储在 `SavePlan::storage_data` 中。
+
+plan (SavePlan) – 使用中的 `SavePlanner` 的本地计划。
+
+调用表示即将进行一次全新的检查点写入。如果用户为此检查点写入设置了 `checkpoint_id`，则可能会存在该 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹/文件的路径，也可以是键值存储的键。
+
+checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。`checkpoint_id` 的含义取决于存储方式。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是一个键。(默认: None)
+
+初始化此实例。
+
+is_coordinator (bool) – 此实例是否负责协调检查点。
+
+返回特定于存储的元数据。这用于在检查点中存储附加信息，这些信息对于提供请求级别的可观察性很有用。`StorageMeta` 在保存调用期间传递给 `SavePlanner`。默认返回 None。
+
+示例:
+
+```python
+from torch.distributed.checkpoint.storage import StorageMeta
+
+class CustomStorageBackend:
+    def get_storage_metadata(self):
+        # 返回将与检查点一起存储的特定于存储的元数据
+        return StorageMeta()
+```
+
+此示例演示了存储后端如何返回 `StorageMeta`
+以将附加元数据附加到检查点。
+
+可选[StorageMeta]
+
+检查给定的 `checkpoint_id` 是否受存储支持。这允许我们启用自动存储选择。
+
+使用规划器解析数据，并从计划中写入所有项目。
+
+子类应针对计划中的每个项目调用 `SavePlanner::resolve_data` 以获取要写入的底层对象的访问权限。
+
+子类应延迟调用 `resolve_data`，因为它可能会分配内存。对于张量，请做出以下假设：
+
+它们可能在任何设备上，包括与 `WriteItem::tensor_data` 上的设备不匹配的情况
+
+它们可能是视图或非连续的。只有投影部分需要被保存。
+
+plan (SavePlan) – 要执行的保存计划。
+
+planner (SavePlanner) – 用于将项目解析为数据的规划器对象。
+
+一个在结束时产生 `WriteResult` 列表的 Future
+
+Future[list[torch.distributed.checkpoint.storage.WriteResult]]
+
+以下类型定义了检查点期间使用的规划器接口：
+
+定义 `load_state_dict` 用来规划加载过程的协议的抽象类。
+
+`LoadPlanner` 是有状态的对象，可用于自定义整个加载过程。
+
+`LoadPlanner` 充当 `state_dict` 的访问代理，因此对它进行的任何转换将对整个过程可见。
+
+在 `load_state_dict` 期间，规划器子类可以预期以下调用序列：
+
+发出加载检查点开始的信号。
+
+处理 `state_dict` 并生成将被发送以进行全局规划的 `LoadPlan`。
+
+获取来自所有 rank 的 `LoadPlan` 并做出任何全局决定。
+
+这对于 `state_dict` 中的每个非张量值调用一次。
+
+它们对于 `state_dict` 中的每个张量值成对调用。
+
+建议用户扩展 `DefaultLoadPlanner` 而不是直接扩展此接口，因为大多数更改都可以通过单个方法的更改来表达。
+
+通常有两种扩展模式：
+
+重写 `state_dict`。这是扩展加载过程的最简单方法，因为它不需要了解 `LoadPlan` 是如何工作的复杂性。我们需要保留对原始 `state_dict` 的引用，因为加载是原位发生的，所以我们必须能够原位执行它
+
+修改 `resolve_tensor` 和 `commit_tensor` 以处理加载时转换。
+
+在 `StorageReader` 完成将数据加载到张量中后调用一次。
+
+提供的张量与调用 `resolve_tensor` 返回的张量相同。仅当此 `LoadPlanner` 需要在将其复制回 `state_dict` 中的张量之前对张量进行后处理时，才需要此方法。
+
+张量的内容将遵循其设备同步模型。
+
+计算全局加载计划并返回每个 rank 的计划。
+
+. 注意：这仅在协调者 rank 上调用
+
+list[torch.distributed.checkpoint.planner.LoadPlan]
+
+根据 `state_dict` 和 `set_up_planner` 提供的元数据创建 `LoadPlan`。
+
+. 注意：这在每个 rank 上都会被调用。
+
+接受来自协调者的计划并返回最终的 `LoadPlan`。
+
+加载由 `read_item`` 和 ``value` 描述的项目。
+
+此方法预期会原位修改底层的 `state_dict`。
+
+`value` 的内容由用于生成正在加载的检查点的 `SavePlanner` 定义。
+
+返回供 `StorageReader` 加载 `read_item` 时使用的 `BytesIO`。
+
+`BytesIO` 应与底层 `state_dict` 中的某个对象别名相同，因为 `StorageReader` 将替换其内容。
+
+返回供 `StorageReader` 加载 `read_item` 时使用的由 `read_item` 描述的张量。
+
+张量应与底层 `state_dict` 中的某个对象别名相同，因为 `StorageReader` 将替换其内容。如果出于任何原因这不可能，规划器可以使用 `commit_tensor` 方法将数据复制回 `state_dict` 中的对象。
+
+初始化此实例以将数据加载到 `state_dict` 中。
+
+. 注意：这在每个 rank 上都会被调用。
+
+定义 `save_state_dict` 用来规划保存过程的协议的抽象类。
+
+`SavePlanner` 是有状态的对象，可用于自定义整个保存过程。
+
+`SavePlanner` 充当 `state_dict` 的访问代理，因此对它进行的任何转换将对整个过程可见。
+
+在 `save_state_dict` 期间，规划器子类可以预期以下调用序列：
+
+发出检查点保存开始的信号。
+
+处理 `state_dict` 并生成将被发送以进行全局规划的 `SavePlan`。
+
+获取来自所有 rank 的 `SavePlan` 并做出任何全局决定。
+
+这使得每个 rank 都有机会根据全局规划决策进行调整。
+
+在 `state_dict` 中查找供存储层写入的值。
+
+建议用户扩展 `DefaultSavePlanner` 而不是直接扩展此接口，因为大多数更改都可以通过单个方法的更改来表达。
+
+通常有 3 种扩展模式：
+
+重写 `state_dict`。这是扩展保存过程的最简单方法，因为它不需要了解 `SavePlan` 是如何工作的复杂性：
+
+串联修改本地计划和查找。这对于需要精细控制数据如何持久化非常有用
+
+利用全局规划步骤做出每个 rank 无法单独做出的集中决策
+
+最后，一些规划器需要在检查点中保存附加的元数据，这是通过让每个 rank 在本地计划中贡献其数据项，然后由全局规划器聚合它们来实现的：
+
+计算全局检查点计划并返回每个 rank 的本地计划。
+
+这仅在协调者 rank 上调用。
+
+tuple[list[torch.distributed.checkpoint.planner.SavePlan], torch.distributed.checkpoint.metadata.Metadata]
+
+计算当前 rank 的保存计划。
+
+这将被聚合并传递给 `create_global_plan`。特定于规划器的数据可以通过 `SavePlan::planner_data` 传递。
+
+这在所有 rank 上调用。
+
+合并 `create_local_plan` 创建的计划和 `create_global_plan` 的结果。
+
+这在所有 rank 上调用。
+
+从 `state_dict` 转换并准备 `write_item` 以供存储，确保幂等性和线程安全。
+
+在 `state_dict` 中查找与 `write_item` 关联的对象，并在存储层使用它之前应用任何转换（例如序列化）。
+
+在每个 rank 上被多次调用，最终 `SavePlan` 中的每个 `WriteItem` 至少调用一次。
+
+此方法应是幂等的且线程安全的。`StorageWriter` 实现可以根据需要自由调用它。
+
+任何分配内存的转换都应在此方法被调用时延迟执行，以减少检查点所需的峰值内存。
+
+返回张量时，它们可以在任何设备或格式上，也可以是视图。存储层有责任弄清楚如何保存它们。
+
+Union[Tensor, BytesIO]
+
+初始化此规划器以保存 `state_dict`。
+
+实现应保存这些值，因为在保存过程的后期将不会提供它们。
+
+这在所有 rank 上调用。
+
+包含有关需要写入存储的内容信息的数据类。
+
+计算底层张量的存储大小，如果这不是张量写入，则为 None。
+
+可选[int] 底层张量的存储大小（以字节为单位，如果有）。
+
+我们提供了一个基于文件系统的存储层：
+
+返回将用于加载检查点的 `checkpoint_id`。
+
+使用文件 IO 的 `StorageWriter` 的基本实现。
+
+此实现做出以下假设和简化：
+
+检查点路径是一个空目录或不存在的目录。
+
+文件创建是原子的
+
+如果启用了 rank 协调，检查点由每个写入请求一个文件加上一个包含序列化元数据的全局 `.metadata` 文件组成。如果未启用 rank 协调，则由一个 rank 本地的 `__{rank}.metadata` 文件包含序列化元数据。
+
+重写 `AsyncStager.stage`
+
+dict[str, Union[~StatefulT, Any]]
+
+我们还提供了其他存储层，包括与 HuggingFace safetensors 交互的存储层：
+
+.. autoclass:: torch.distributed.checkpoint.HuggingFaceStorageReader :members:
+
+.. autoclass:: torch.distributed.checkpoint.HuggingFaceStorageWriter :members:
+
+.. autoclass:: torch.distributed.checkpoint.QuantizedHuggingFaceStorageReader :members:
+
+我们提供了 `LoadPlanner` 和 `SavePlanner` 的默认实现，它们可以处理所有 `torch.distributed` 构造，例如 FSDP、DDP、`ShardedTensor` 和 `DistributedTensor`。
+
+从规划器接口扩展，以便轻松扩展默认规划器。
+
+从规划器接口扩展，以便轻松扩展默认规划器。
+
+`DefaultLoadPlanner` 在 `LoadPlanner` 之上添加了多项功能。
+
+特别是它添加了以下内容：
+
+flatten_state_dict：处理带有嵌套字典的 state_dict flatten_sharded_tensors：对于二维并行模式下的 FSDP allow_partial_load：如果为 False，当某个键存在于 `state_dict` 中但不在检查点中时，将引发运行时错误。
+
+从规划器接口扩展，以便轻松扩展默认规划器。
+
+从规划器接口扩展，以便轻松扩展默认规划器。
+
+由于早期设计决策，即使原始的未并行化模型完全相同，FSDP 和 DDP 的状态字典可能具有不同的键或全限定名称（例如 `layer1.weight`）。此外，FSDP 提供各种类型的模型状态字典，例如完整和分片状态字典。此外，优化器状态字典使用参数 ID 而不是全限定名称来标识参数，这在使用并行化（例如流水线并行）时可能会导致问题。
+
+为了应对这些挑战，我们提供了一套 API，让用户能够轻松管理 state_dict。`get_model_state_dict()` 返回一个模型状态字典，其键与未并行化模型状态字典返回的键一致。类似地，`get_optimizer_state_dict()` 提供的应用于所有并行化中一致的键的优化器状态字典。为了实现这种一致性，`get_optimizer_state_dict()` 会将参数 ID 转换为与未并行化模型状态字典中完全相同的全限定名称。
+
+请注意，这些 API 返回的结果可以直接与 `torch.distributed.checkpoint.save()` 和 `torch.distributed.checkpoint.load()` 方法一起使用，无需任何额外转换。
+
+提供 `set_model_state_dict()` 和 `set_optimizer_state_dict()` 用于加载由各自获取 API 生成的模型和优化器 state_dict。
+
+请注意，只能在优化器上调用 `backward()` 之前或调用 `step()` 之后调用 `set_optimizer_state_dict()`。
+
+请注意，此功能是实验性的，API 签名将来可能会发生变化。
+
+返回模型 state_dict 和优化器 state_dict。
+
+`get_state_dict` 可以处理任何由 PyTorch FSDP/fully_shard、DDP/replicate、tensor_parallel/parallelize_module 以及这些并行化的任意组合进行并行化的模块。`get_state_dict` 的主要功能是：1.) 返回可以使用不同数量的训练器和/或不同的并行化进行重新分片的模型和优化器 state_dict。2.) 隐藏特定于并行化的 state_dict API。用户不必调用这些 API。3.) 对结果 state_dict 进行健全性检查。
+
+结果状态字典的键是标准的 FQN（全限定名称）。标准 FQN 指的是基于参数在 `nn.Module` 层次结构中位置的 FQN。更具体地说，当模块未被任何并行化分布时，参数的标准 FQN 是 `module.named_parameters()` 或 `module.named_buffers()` 返回的 FQN。由于优化器内部使用参数 ID 来表示参数，因此在调用此 API 时会有一个从参数 ID 到标准 FQN 的转换。
+
+`get_state_dict` 也可以处理未并行化的模块。在这种情况下，`get_state_dict` 仅执行一项功能 —— 将优化器参数 ID 转换为标准 FQN。
+
+model (nn.Module) – 模型的 nn.Module。
+
+optimizers (Union[None, Optimizer, Iterable[Optimizer]]) – 用于优化 `model` 的优化器。
+
+submodules (已弃用) – 可选[set[nn.Module]]: 仅返回属于这些子模块的模型参数。
+
+options (StateDictOptions) – 控制如何返回模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+包含模型 state_dict 和优化器 state_dict 的元组。
+
+Tuple[Dict[str, ValueType], OptimizerStateType]
+
+返回模型的模型 state_dict。
+
+详情请参见 `get_state_dict` 的用法。
+
+model (nn.Module) – 模型的 nn.Module。
+
+submodules (已弃用) – 可选[set[nn.Module]]: 仅返回属于这些子模块的模型参数。
+
+options (StateDictOptions) – 控制如何返回模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+模型的 state_dict。
+
+返回优化器的组合 state_dict。
+
+详情请参见 `get_state_dict` 的用法。
+
+model (nn.Module) – 模型的 nn.Module。
+
+optimizers (Union[None, Optimizer, Iterable[Optimizer]]) – 用于优化 `model` 的优化器。
+
+submodules (已弃用) – 可选[set[nn.Module]]: 仅返回属于这些子模块的模型参数。
+
+options (StateDictOptions) – 控制如何返回模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+优化器的 state_dict。
+
+加载模型 state_dict 和优化器 state_dict。
+
+`get_state_dict` 的对应方法，用于将 state_dict 设置到模型和优化器中。给定的 `model_state_dict` 和 `optim_state_dict` 不必由 `get_state_dict` 返回，但必须满足以下要求：1) 所有 FQN 都是 `get_state_dict` 中定义的标准 FQN，2) 如果张量被分片，它必须是 `ShardedTensor` 或 `DTensor`，3) 优化器 state_dict 不能包含参数 ID；键应该是标准 FQN。
+
+在优化器上被调用。否则，优化器状态将无法正确初始化。
+
+model (nn.Module) – 模型的 nn.Module。
+
+optimizers (Union[Optimizer, Iterable[Optimizer]]) – 用于优化 `model` 的优化器。
+
+model_state_dict (Dict[str, ValueType]) – (Union[Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]]): 要加载的模型 state_dict。如果 `model_state_dict` 的键是 `nn.Module`，则该键是 `model` 的子模块，值应为该子模块的 state_dict。加载 state_dict 时，子模块的前缀将被附加到 state_dict 中。
+
+optim_state_dict (OptimizerStateType) – OptimizerStateType: 要加载的优化器 state_dict。
+
+options (StateDictOptions) – 控制如何加载模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+`missing_keys` 是一个字符串列表，包含模型 state_dict 中缺失的键。`unexpected_keys` 是一个字符串列表，包含模型 state_dict 中意外的键。
+
+`missing_keys` 是一个字符串列表，包含模型 state_dict 中缺失的键。
+
+`unexpected_keys` 是一个字符串列表，包含模型 state_dict 中意外的键。
+
+包含 `missing_keys` 和 `unexpected_keys` 字段的 NamedTuple
+
+加载模型 state_dict。
+
+`get_model_state_dict` 的对应方法，用于将 state_dict 设置到模型中。详情请参见 `set_state_dict` 的用法。
+
+model (nn.Module) – 模型的 nn.Module。
+
+model_state_dict (Dict[str, ValueType]) – (Dict[str, ValueType]): 要加载的模型 state_dict。如果 `model_state_dict` 的键是 `nn.Module`，则该键是 `model` 的子模块，值应为该子模块的 state_dict。加载 state_dict 时，子模块的前缀将被附加到 state_dict 中。
+
+options (StateDictOptions) – 控制如何加载模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+`missing_keys` 是一个字符串列表，包含缺失的键 `unexpected_keys` 是一个字符串列表，包含意外的键
+
+`missing_keys` 是一个字符串列表，包含缺失的键
+
+`unexpected_keys` 是一个字符串列表，包含意外的键
+
+包含 `missing_keys` 和 `unexpected_keys` 字段的 NamedTuple
+
+加载优化器 state_dict。
+
+`get_optimizer_state_dict` 的对应方法，用于将 state_dict 设置到优化器中。详情请参见 `set_state_dict` 的用法。
+
+在优化器上调用 `step()`。否则，优化器状态将无法正确初始化。
+
+model (nn.Module) – 模型的 nn.Module。
+
+optimizers (Union[Optimizer, Iterable[Optimizer]]) – 用于优化 `model` 的优化器。
+
+optim_state_dict (OptimizerStateType) – OptimizerStateType: 要加载的优化器 state_dict。
+
+options (StateDictOptions) – 控制如何加载模型 state_dict 和优化器 state_dict 的选项。详情请参见 `StateDictOptions`。
+
+此数据类指定 `get_state_dict`/`set_state_dict` 将如何工作。
+
+full_state_dict: 如果设置为 True，则返回的 state_dict 中的所有张量将被收集。返回的 state_dict 中将不存在 ShardedTensor 和 DTensor。
+
+cpu_offload: 将所有张量卸载到 CPU。为了防止 CPU 发生 OOM，如果 `full_state_dict` 也为 true，则只有 rank0 将获取 state_dict，而所有其他 rank 将获取空的 state_dict。
+
+ignore_frozen_params: 如果值为 True，则返回的 state_dict 将不包含任何冻结的参数 —— 即 `requires_grad` 为 False 的参数。默认值为 False。
+
+keep_submodule_prefixes (已弃用): 当 `submodules` 不为 None 时，此选项指示是否保留 state_dict 键中的子模块前缀。例如，如果子模块是 `module.pretrain` 并且参数的完整 FQN 是参数的 `pretrain.layer1.weight`。当此选项为 True 时，返回的 state_dict 中参数的键将是 `pretrain.layer1.weight`。如果选项为 False，则键将是 `layer1.weight`。请注意，如果 `keep_submodule_prefixes` 为 False，可能会出现冲突的 FQN，因此 `submodules` 中应该只有一个子模块。
+
+strict: `set_state_dict` 调用 `model.load_state_dict()` 时的严格选项。
+
+完整的 state_dict，并会将 state_dict/optim_state_dict 中的张量逐一广播给其他 rank。其他 rank 将接收张量并根据模型和优化器中的本地分片进行分片。使用此选项时必须将 `full_state_dict` 设置为 True。此选项目前仅支持 DTensor，不支持遗留的 ShardedTensor。
+
+对于习惯使用和共享 `torch.save` 格式的模型的用户，提供了以下方法，它们提供用于在格式之间进行转换的离线实用程序。
+
+给定一个
+
+## torch.distributed.tensor#
+
+**URL:** https://pytorch.org/docs/stable/distributed.tensor.html
+
+**目录：**
+- torch.distributed.tensor#
+- PyTorch DTensor (Distributed Tensor)#
+  - DTensor Class APIs#
+  - DeviceMesh as the distributed communicator#
+  - DTensor Placement Types#
+- Different ways to create a DTensor#
+  - Create DTensor from a logical torch.Tensor#
+  - DTensor Factory Functions#
+  - Random Operations#
+- Debugging#
+
+创建时间：2025 年 6 月 13 日 | 最后更新时间：2025 年 8 月 23 日
+
+torch.distributed.tensor 目前处于 alpha 状态并在开发中，我们承诺对文档中列出的大多数 API 保持向后兼容，但如有必要，可能会发生 API 变更。
+
+PyTorch DTensor 提供了简单灵活的张量切分原语，可透明地处理分布式逻辑，包括分片存储、算子计算以及跨设备/主机的集合通信。DTensor 可用于构建不同的并行解决方案，并支持在处理多维分片时的分片 state_dict 表示。
+
+请查看基于 DTensor 构建的 PyTorch 原生并行解决方案的示例：
+
+DTensor 遵循 SPMD（单程序多数据）编程模型，使用户能够像编写单设备程序一样编写分布式程序，并具有相同的收敛属性。它通过指定 DeviceMesh 和 Placement 来提供统一的张量分片布局（DTensor Layout）：
+
+DeviceMesh 使用 n 维数组表示集群的设备拓扑和通信器。
+
+Placement 描述逻辑张量在 DeviceMesh 上的分片布局。DTensor 支持三种类型的放置方式：Shard（分片）、Replicate（复制）和 Partial（部分）。
+
+DTensor 是 torch.Tensor 的子类。这意味着一旦创建了 DTensor，就可以像使用 torch.Tensor 一样使用它，包括像在单设备上运行一样运行不同类型的 PyTorch 算子，从而允许 PyTorch 算子进行正确的分布式计算。
+
+除了现有的 torch.Tensor 方法外，它还提供了一组附加方法来与 torch.Tensor 交互、将 DTensor Layout 重新分配为新的 DTensor、在所有设备上获取完整的张量内容等。
+
+DTensor (Distributed Tensor) 是 torch.Tensor 的子类，它提供类似单设备的抽象，以便对多设备 torch.Tensor 进行编程。它通过 DeviceMesh 和以下类型的 Placement 来描述分布式张量分片布局（DTensor Layout）：
+
+Shard：在 DeviceMesh 维度对应的设备上，对张量的维度 dim 进行切分
+
+Replicate：在 DeviceMesh 维度的设备上复制张量
+
+Partial：在 DeviceMesh 维度的设备上，张量处于待归约状态
+
+在调用 PyTorch 算子时，DTensor 会重写 PyTorch 算子以执行分片计算，并在必要时发起通信。伴随着算子计算，DTensor 会（根据算子本身的语义）正确地转换或传播放置方式（DTensor Layout），并生成新的 DTensor 输出。
+
+为了确保在调用 PyTorch 算子时 DTensor 分片计算的数值正确性，DTensor 要求该算子的每个张量参数都必须是 DTensor。
+
+在此直接使用 Tensor 子类构造函数并不是创建 DTensor 的推荐方法（即它不能正确处理自动求导，因此不是公开的 API）。请参阅 create_dtensor 部分以了解如何创建 DTensor。
+
+返回一个 ChunkStorageMetadata 列表，它是一个描述当前 rank 上本地分片/副本的大小/偏移量的数据类。对于 DTensor，每个 rank 将有一个本地分片/副本，因此返回的列表通常只有一个元素。
+
+此 dunder 方法主要用于分布式检查点目的。
+
+一个表示当前 rank 上分片大小/偏移量的 List[ChunkStorageMetadata] 对象。
+
+根据指定的 device_mesh 和 placements，在每个 rank 上通过本地 torch.Tensor 创建一个 DTensor。
+
+local_tensor (torch.Tensor) – 每个 rank 上的本地 torch.Tensor。
+
+device_mesh (DeviceMesh, 可选) – 用于放置张量的 DeviceMesh，如果未指定，必须在 DeviceMesh 上下文管理器下调用，默认值：None
+
+placements (List[Placement], 可选) – 描述如何将本地 torch.Tensor 放置到 DeviceMesh 上的放置方式，其元素数量必须与 device_mesh.ndim 相同。
+
+run_check (bool, 可选) – 以额外的通信为代价，跨 rank 执行完整性检查，检查每个本地张量的元信息以确保正确性。如果 placements 中包含 Replicate，设备网格维度上第一个 rank 的数据将被广播到其他 rank。默认值：False
+
+shape (torch.Size, 可选) – 一个整数列表，指定建立在 local_tensor 之上的 DTensor 的大小。注意，如果各个 rank 上的 local_tensor 形状不同，则需要提供此参数。如果未提供，将假设给定的分布式张量在各个 rank 上被均匀切分来计算形状。默认值：None
+
+stride (tuple, 可选) – 一个整数列表，指定 DTensor 的步长。如果未提供，将假设给定的分布式张量在各个 rank 上被均匀切分来计算步长。默认值：None
+
+当 run_check=False 时，用户有责任确保跨 rank 传入的本地张量是正确的（即对于 Shard(dim) 放置方式，张量应被分片；对于 Replicate() 放置方式，张量应被复制）。否则，创建的 DTensor 的行为是未定义的。
+
+from_local 是可微的，创建的 DTensor 对象的 requires_grad 将取决于 local_tensor 是否需要梯度。
+
+返回此 DTensor 的完整张量。它将执行必要的集合通信，从其 DeviceMesh 上的其他 rank 收集本地张量并将它们拼接在一起。它是以下代码的语法糖：
+
+dtensor.redistribute(placements=[Replicate()] * mesh.ndim).to_local()
+
+grad_placements (List[Placement], 可选) – 该放置方式描述了此函数返回的完整张量的任何梯度布局的未来布局。full_tensor 将 DTensor 转换为完整的 torch.Tensor，并且返回的 torch.tensor 稍后在代码中可能不会用作原始复制的 DTensor 布局。此参数是用户可以在返回张量的梯度布局与原始复制 DTensor 布局不匹配时提供给自动求导的提示。如果未指定，我们将假设完整张量的梯度布局为复制。
+
+一个表示此 DTensor 完整张量的 torch.Tensor 对象。
+
+full_tensor 是可微的。
+
+redistribute 执行必要的集合操作，将当前 DTensor 从其当前的 placements 重新分配为新的 placements，或者从其当前的 DeviceMesh 重新分配到新的 DeviceMesh。即，我们可以通过为 DeviceMesh 的每个维度指定 Replicate 放置方式，将分片的 DTensor 转换为复制的 DTensor。
+
+在一个设备网格维度上从当前 placements 重新分配到新的 placements 时，我们将执行以下包含通信集合操作或本地操作：
+
+Shard(dim) -> Replicate(): all_gather
+
+Shard(src_dim) -> Shard(dst_dim): all_to_all
+
+Replicate() -> Shard(dim): 本地分块（即 torch.chunk）
+
+Partial() -> Replicate(): all_reduce
+
+Partial() -> Shard(dim): reduce_scatter
+
+redistribute 能够正确地为在 1 维或 N 维 DeviceMesh 上创建的 DTensor 找出必要的重新分配步骤。
+
+device_mesh (DeviceMesh, 可选) – 用于放置 DTensor 的 DeviceMesh。如果未指定，它将使用当前 DTensor 的 DeviceMesh。默认值：None
+
+placements (List[Placement], 可选) – 描述如何将 DTensor 放置到 DeviceMesh 中的新放置方式，其元素数量必须与 device_mesh.ndim 相同。默认值：在所有网格维度上进行复制
+
+async_op (bool, 可选) – 是否异步执行 DTensor 重新分配操作。默认: False
+
+forward_dtype (torch.dtype, 可选) – 在其前向传播中重新分配本地张量之前，可以将本地张量数据类型转换为 forward_dtype。生成的 DTensor 将处于 forward_dtype 默认: None。
+
+backward_dtype (torch.dtype, 可选) – 在其反向传播中重新分配本地张量之前，可以将本地张量数据类型转换为 backward_dtype。生成的 DTensor 梯度将被转换回当前 DTensor 的数据类型。默认: None
+
+redistribute 是可微的，这意味着用户不需要担心 redistribute 操作的反向传播公式。
+
+redistribute 目前仅支持在同一个 DeviceMesh 上重新分配 DTensor，如果您需要将 DTensor 重新分配到不同的 DeviceMesh，请提交一个 issue。
+
+获取此 DTensor 在其当前 rank 上的本地张量。对于分片，它返回逻辑张量视图的一个本地分片；对于复制，它返回其当前 rank 上的副本。
+
+grad_placements (List[Placement], 可选) – 该放置方式描述了此函数返回的张量的任何梯度布局的未来布局。to_local 将 DTensor 转换为本地张量，并且返回的本地张量稍后在代码中可能不会用作原始 DTensor 布局。此参数是用户可以在返回张量的梯度布局与原始 DTensor 布局不匹配时提供给自动求导的提示。如果未指定，我们将假设梯度布局保持与原始 DTensor 相同，并将其用于梯度计算。
+
+一个 torch.Tensor 或 AsyncCollectiveTensor 对象。它表示其当前 rank 上的本地张量。当返回 AsyncCollectiveTensor 对象时，意味着本地张量尚未就绪（即通信尚未完成）。在这种情况下，用户需要调用 wait 等待本地张量就绪。
+
+to_local 是可微的，返回的本地张量的 requires_grad 将取决于该 DTensor 是否需要梯度。
+
+与此 DTensor 对象关联的 DeviceMesh 属性。
+
+device_mesh 是只读属性，不能被设置。
+
+此 DTensor 的 placements 属性，描述了该 DTensor 在其 DeviceMesh 上的布局。
+
+placements 是只读属性，不能被设置。
+
+DeviceMesh 从 DTensor 构建而来，作为描述集群设备拓扑并表示多维通信器（构建在 ProcessGroup 之上）的抽象。要详细了解如何创建/使用 DeviceMesh，请参阅 DeviceMesh 指南。
+
+DTensor 在每个 DeviceMesh 维度上支持以下类型的 Placement：
+
+Shard(dim) 放置方式描述了 DTensor 在对应的 DeviceMesh 维度上对张量维度 dim 的切分，其中 DeviceMesh 维度上的每个 rank 仅持有全局张量的一个分片/部分。Shard(dim) 放置方式遵循 torch.chunk(dim) 语义，其中当张量维度无法在 DeviceMesh 维度上被整除时，DeviceMesh 维度上的最后几个分片可能为空。Shard 放置方式可以被所有的 DTensor API 使用（即 distribute_tensor、from_local 等）
+
+dim (int) – 描述 DTensor 在其对应的 DeviceMesh 维度上进行切分的张量维度。
+
+当张量维度的大小无法在 DeviceMesh 维度上整除时，对该张量维度进行分片目前是实验性的，可能会发生变化。
+
+Replicate() 放置方式描述了 DTensor 在对应的 DeviceMesh 维度上进行复制，其中 DeviceMesh 维度上的每个 rank 都持有全局张量的一个副本。Replicate 放置方式可以被所有的 DTensor API 使用（即 distribute_tensor、DTensor.from_local 等）
+
+Partial(reduce_op) 放置方式描述了在指定的 DeviceMesh 维度上处于待归约状态的 DTensor，其中 DeviceMesh 维度上的每个 rank 都持有全局张量的部分值。用户可以使用 redistribute 将 Partial DTensor 重新分配为指定 DeviceMesh 维度上的 Replicate 或 Shard(dim) 放置方式，这将在底层触发必要的通信操作（即 allreduce、reduce_scatter）。
+
+reduce_op (str, 可选) – 用于部分 DTensor 生成复制/分片 DTensor 的归约操作。仅支持逐元素归约操作，包括："sum"、"avg"、"product"、"max"、"min"，默认值："sum"。
+
+Partial 放置方式可以作为 DTensor 算子的结果生成，并且只能被 DTensor.from_local API 使用。
+
+Placement 类型的基类，它描述了 DTensor 是如何放置到 DeviceMesh 上的。Placement 和 DeviceMesh 一起可以描述 DTensor Layout。它是三种主要 DTensor Placement 类型的基类：Shard、Replicate 和 Partial。
+
+此类不适合直接使用，主要作为类型提示存根。
+
+distribute_tensor() 在每个 rank 上通过逻辑或“全局” torch.Tensor 创建一个 DTensor。这可用于对叶子 torch.Tensor（即模型参数/缓冲区和输入）进行切分。
+
+DTensor.from_local() 在每个 rank 上通过本地 torch.Tensor 创建一个 DTensor，这可用于从非叶子 torch.Tensor（即前向/反向传播过程中的中间激活张量）创建 DTensor。
+
+DTensor 提供了专用的张量工厂函数（例如 empty()、ones()、randn() 等），允许通过直接指定 DeviceMesh 和 Placement 进行不同的 DTensor 创建。与 distribute_tensor() 相比，这可以直接在设备上实例化分片内存，而不是在初始化逻辑张量内存之后再执行分片。
+
+torch.distributed 中的 SPMD（单程序多数据）编程模型会启动多个进程（例如通过 torchrun）来执行同一个程序，这意味着程序中的模型将首先在不同的进程上被初始化（即模型可能在 CPU、或 meta 设备上初始化，或者在内存充足的情况下直接在 GPU 上初始化）。
+
+DTensor 提供了一个 distribute_tensor() API，可以将模型权重或张量切分为多个 DTensor，它会在每个进程上通过“逻辑”张量创建一个 DTensor。这将使创建的 DTensor 遵循单设备语义，这对于数值正确性至关重要。
+
+根据指定的 placements 将叶子 torch.Tensor（即 nn.Parameter/缓冲区）分发到 device_mesh。device_mesh 和 placements 的秩必须相同。要分发的张量是逻辑或“全局”张量，并且该 API 将使用 DeviceMesh 维度上第一个 rank 的张量作为真实数据源，以保留单设备语义。如果您想在自动求导计算过程中构建 DTensor，请改用 DTensor.from_local()。
+
+tensor (torch.Tensor) – 要分发的 torch.Tensor。请注意，如果您想在一个不能被该网格维度中的设备数整除的维度上对张量进行切分，我们将使用 torch.chunk 语义对张量进行切分并分散这些分片。这种不均匀分片的行为是实验性的，可能会发生变化。
+
+device_mesh (DeviceMesh, 可选) – 用于分发张量的 DeviceMesh，如果未指定，必须在 DeviceMesh 上下文管理器下调用，默认值：None
+
+placements (List[Placement], 可选) – 描述如何将张量放置到 DeviceMesh 上的放置方式，其元素数量必须与 device_mesh.ndim 相同。如果未指定，我们默认将跨越 device_mesh 从每个维度的第一个 rank 复制张量。
+
+src_data_rank (int, 可选) – 逻辑/全局张量的源数据 rank，distribute_tensor() 使用它将分片/副本分散/广播到其他 rank。默认情况下，我们在每个 DeviceMesh 维度上使用 group_rank=0 作为源数据，以保留单设备语义。如果显式传入 None，distribute_tensor() 将仅使用其本地数据，而不是尝试通过分散/广播来保留单设备语义。默认: 0
+
+一个 DTensor 或 XLAShardedTensor 对象。
+
+当使用 xla 设备类型初始化 DeviceMesh 时，distribute_tensor 返回的是 XLAShardedTensor。有关更多详细信息，请参见此 issue。XLA 集成是实验性的，可能会发生变化。
+
+除了 distribute_tensor() 之外，DTensor 还提供了一个 distribute_module() API，以便更容易在 nn.Module 级别上进行分片。
+
+该函数公开了三个函数来控制模块的参数/输入/输出：
+
+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，默认情况下我们将在网格上复制该模块的所有模块参数。
+
+input_fn (Callable) – 指定输入分布，例如可以控制模块的输入如何被分片。input_fn 将作为模块的 forward_pre_hook（前向前置钩子）安装。
+
+output_fn (Callable) – 指定输出分布，例如可以控制输出如何被分片，或将其转换回 torch.Tensor。output_fn 将作为模块的 forward_hook（前向后置钩子）安装。
+
+一个包含全为 DTensor 参数/缓冲区的模块。
+
+当使用 xla 设备类型初始化 DeviceMesh 时，distribute_module 返回带有 PyTorch/XLA SPMD 注解参数的 nn.Module。有关更多详细信息，请参见此 issue。XLA 集成是实验性的，可能会发生变化。
+
+DTensor 还提供了专用的张量工厂函数，允许通过额外指定所创建 DTensor 的 DeviceMesh 和 Placement，直接使用类似 torch.Tensor 的工厂函数 API（即 torch.ones、torch.empty 等）来创建 DTensor：
+
+返回一个填充有标量值 0 的 DTensor。
+
+size (int...) – 定义输出 DTensor 形状的整数序列。可以是可变数量的参数或像列表或元组这样的集合。例如：zeros(1,2,3..) 或 zeros([1,2,3..]) 或 zeros((1,2,3..))
+
+requires_grad (bool, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+dtype (torch.dtype, 可选) – 返回的 DTensor 所需的数据类型。默认: 如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。
+
+layout (torch.layout, 可选) – 返回的 DTensor 所需的布局。默认: torch.strided。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+返回一个填充有标量值 1 的 DTensor，其形状由可变参数 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, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+返回一个用未初始化数据填充的 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, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+根据 device_mesh 和 placements 返回一个用 fill_value 填充的 DTensor，其形状由参数 size 定义。
+
+size (int...) – 定义输出 DTensor 形状的整数序列。可以是可变数量的参数或像列表或元组这样的集合。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))
+
+fill_value (Scalar) – 用于填充输出张量的值。
+
+dtype (torch.dtype, 可选) – 返回的 DTensor 所需的数据类型。默认: 如果为 None，则使用全局默认值（参见 torch.set_default_dtype()）。
+
+layout (torch.layout, 可选) – 返回的 DTensor 所需的布局。默认: torch.strided。
+
+requires_grad (bool, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+返回一个用 [0, 1) 区间内均匀分布的随机数填充的 DTensor。张量的形状由可变参数 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, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+返回一个用均值为 0、方差为 1 的正态分布中提取的随机数填充的 DTensor。张量的形状由可变参数 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, 可选) – 自动求导是否应记录在返回的 DTensor 上的操作。默认: False。
+
+device_mesh – DeviceMesh 类型，包含 rank 的网格信息。
+
+placements – Placement 类型的序列：Shard、Replicate。
+
+每个 rank 上的一个 DTensor 对象。
+
+DTensor 提供了分布式 RNG 功能，以确保对分片张量执行的随机操作能获得唯一的值，而对复制张量执行的随机操作能获得相同的值。该系统要求所有参与的 rank（例如 SPMD rank）在执行每次 dtensor 随机操作之前，都以相同的生成器状态开始；如果满足此条件，它将确保在每次 dtensor 随机操作完成后，它们最终都达到相同的状态。在随机操作期间不执行任何通信来同步 RNG 状态。
+
+接受 generator 关键字参数的算子将使用用户传入的生成器（如果传入了的话），否则使用该设备的默认生成器。无论使用哪种生成器，它都会在 DTensor 操作之后更新。将同一个生成器同时用于 DTensor 和非 DTensor 操作是有效的，但在此情况下，必须注意确保非 DTensor 操作在所有 rank 上同等地更新生成器状态。
+
+当将 DTensor 与流水线并行结合使用时，每个流水线阶段的 rank 应使用不同的种子，而流水线阶段内的 rank 应使用相同的种子。
+
+DTensor 的 RNG 基础设施基于 philox 的 RNG 算法，并支持任何基于 philox 的后端（cuda 及其他类 cuda 设备），但遗憾的是尚不支持 CPU 后端。
+
+启动程序时，您可以使用来自 torch._logging 的 TORCH_LOGS 环境变量开启额外的日志记录：
+
+TORCH_LOGS=+dtensor 将显示 logging.DEBUG 消息及以上的所有级别。
+
+TORCH_LOGS=dtensor 将显示 logging.INFO 及以上的消息。
+
+TORCH_LOGS=-dtensor 将显示 logging.WARNING 及以上的消息。
+
+为了调试应用了 DTensor 的程序，并了解底层发生了哪些集合通信的更多细节，DTensor 提供了 CommDebugMode：
+
+CommDebugMode 是一个上下文管理器，用于统计其上下文内的功能集合通信数量。它是通过 TorchDispatchMode 来实现这一点的。
+
+目前尚未支持所有的集合通信。
+
+生成显示模块级操作和集合通信追踪信息的详细表格。信息量取决于 noise_level。
+
+打印模块级别的集合通信计数
+
+打印不包含在平凡操作中的 dTensor 操作，以及模块信息
+
+打印不包含在平凡操作中的操作
+
+打印所有操作
+
+创建用于构建浏览器可视化的 json 文件。0. 打印模块级别的集合通信计数 1. 打印不包含在平凡操作中的 dTensor 操作 2. 打印不包含在平凡操作中的操作 3. 打印所有操作。
+
+以字典形式返回通信计数。
+
+以字典形式表示的通信计数。
+
+dict[str, dict[str, Any]]
+
+dict[str, dict[str, Any]]
+
+作为控制台 CommDebugMode 输出的替代方案，将内容写入由用户指定的文件中。
+
+为了可视化维度少于 3 维的 DTensor 的分片情况，DTensor 提供了 visualize_sharding()：
+
+在终端中可视化 1 维或 2 维 DTensor 的分片。
+
+此操作需要安装 tabulate 包，或者 rich 和 matplotlib 包。对于空张量，不会打印任何分片信息。
+
+DTensor 还提供了一系列实验性功能。这些功能要么处于原型阶段，要么基本功能已经完成但正在寻找用户反馈。如果您对这些功能有反馈意见，请向 PyTorch 提交 issue。
+
+context_parallel 是一个用于启用上下文并行 (CP) 的实验性 API。此 API 执行两个操作：1）使用支持 CP 的版本来修补 SDPA (torch.nn.functional.scaled_dot_product_attention)，2）沿序列维度对缓冲区进行分片，每个 rank 将根据网格保留相应的分片。
+
+mesh (DeviceMesh) – 用于上下文并行的设备网格。
+
+buffers (可选[List[torch.Tensor]]) – 其使用依赖于序列维度的缓冲区。示例包括输入批次、标签和位置嵌入缓冲区。这些缓冲区必须沿序列维度进行分片以确保准确性。分片将就地发生，缓冲区的形状将在上下文内改变。上下文结束后缓冲区将被恢复。no_restore_buffers 可用于指定哪些缓冲区不需要恢复。请注意，buffers 不应包含任何 nn.Parameter。
+
+buffer_seq_dims (可选[List[int]]) – 缓冲区的序列维度。
+
+no_restore_buffers (可选[Set[torch.Tensor]]) – 此集合中的缓冲区在上下文退出后将不会被恢复。此集合必须是 buffers 的子集。如果缓冲区在上下文退出后不再使用，可以将这些缓冲区放入此列表中以避免额外的恢复时间。
+
+Generator[None, None, None]
+
+torch.distributed.tensor.experimental.context_parallel 是 PyTorch 中的一个原型功能。该 API 可能会发生变化。
+
+local_map() 是一个实验性 API，允许用户将 DTensor 传递给为应用于 torch.Tensor 而编写的函数。它的实现方式是提取 DTensor 的本地组件，调用该函数，然后根据 out_placements 将输出包装回 DTensor。
+
+func (Callable) – 应用于每个 DTensor 本地分片的函数。
+
+out_placements (Union[PlacementType, Tuple[PlacementType, …]]) – func 展平输出中 DTensor 所需的放置方式。如果展平的输出是单个值，则 out_placements 应为 PlacementType 类型。如果展平的输出包含多个值，则 out_placements 应为 PlacementType 值的元组，并与展平的输出 1:1 映射。此外，对于 Tensor 输出，我们使用 PlacementType 作为其放置方式（即 Tuple[Placement] 值）。对于非 Tensor 输出，PlacementType 应为 None。请注意，唯一的例外是没有传入 DTensor 参数的情况。在这种情况下，即使 out_placements 不为 None，结果函数也应忽略所需的放置方式，因为该函数没有使用 DTensor 运行。
+
+in_placements (Tuple[PlacementType, …], 可选) – func 的展平输入中 DTensor 所需的放置方式。如果指定了 in_placements，local_map() 将检查每个 DTensor 参数的放置方式是否与所需的放置方式相同。如果放置方式不同并且 redistribute_inputs 为 False，则会引发异常。如果 redistribute_inputs 为 True，则在将其本地张量传递给 func 之前，将首先把参数重新分配为所需的分片放置方式。唯一的例外是当所需的放置方式不为 None 且参数是 torch.Tensor 时。在这种情况下，将跳过放置方式检查，并直接将参数传递给 func。如果 in_placements 为 None，则不会执行任何放置方式检查。默认: None
+
+in_grad_placements (Tuple[PlacementType, …], 可选) – 与展平输入 DTensor 对应的 DTensor 梯度的放置提示。此参数是用户可以提供给 to_local() 的提示，以防本地张量输入的梯度布局与其 DTensor 输入布局不匹配。如果未指定，我们将假设本地张量输入的梯度布局保持与原始 DTensor 输入相同，并将其用于梯度计算。默认: None。
+
+device_mesh (DeviceMesh, 可选) – 输出 DTensor 放置的设备网格。如果未指定，这将从第一个输入 DTensor 的设备网格中推断出来。默认: None。
+
+redistribute_inputs (bool, 可选) –指示当输入 DTensor 的放置方式与所需的输入放置方式不同时是否对其进行重新分片的布尔值。如果此值为 False 并且某个 DTensor 输入具有不同的放置方式，则会引发异常。默认: False。
+
+一个将 func 应用于输入 DTensor 的每个本地分片并返回由 func 的返回值构建的 DTensor 的可调用对象。
+
+AssertionError – 对于任何非 DTensor 的输出，我们要求其在 out_placements 中对应的输出放置方式为 None。如果不是这种情况，将引发 AssertionError。
+
+ValueError – 如果 redistribute_inputs=False，但输入 DTensor 根据 in_placements 需要进行重新分配。
+
+此 API 目前是实验性的，可能会发生变化。
+
+register_sharding() 是一个实验性 API，允许用户在张量输入和输出为 DTensor 时为算子注册分片策略。它可能在以下情况非常有用：（1）当算子不存在默认的分片策略时，例如当算子是 DTensor 尚不支持的自定义算子时；（2）当用户希望覆盖现有算子的默认分片策略时。
+
+op (Union[OpOverload, List[OpOverload]]) – 用于注册自定义分片函数的一个算子或算子列表。
+
+一个函数装饰器，可用于包装一个定义了 op 中指定的算子分片策略的函数。定义的分片策略将被注册到 DTensor，并且如果 DTensor 已经实现了该算子，它将覆盖默认的分片策略。自定义分片函数接收与原始算子相同的输入（除非某个参数是 torch.Tensor，它将被替换为 DTensor 内部使用的类张量对象）。该函数应返回一个由 2 元组组成的序列，每个元组指定可接受的输出放置方式及其对应的输入放置方式。
+
+此 API 目前是实验性的，可能会发生变化
+
+## FullyShardedDataParallel#
+
+**URL:** https://pytorch.org/docs/stable/fsdp.html
+
+**目录：**
+- FullyShardedDataParallel#
+
+创建时间：2022年2月2日 | 最后更新时间：2025年6月11日
+
+一个用于在数据并行工作进程间对模块参数进行分片的包装器。
+
+这受到了 Xu 等人以及 DeepSpeed 的 ZeRO Stage 3 的启发。FullyShardedDataParallel 通常被简称为 FSDP。
+
+使用 FSDP 涉及包装你的模块，然后在之后初始化优化器。这是必需的，因为 FSDP 会更改参数变量。
+
+在设置 FSDP 时，你需要考虑目标 CUDA 设备。如果设备具有 ID (dev_id)，你有三个选项：
+
+将该模块放置在该设备上
+
+使用 `torch.cuda.set_device(dev_id)` 设置设备
+
+将 dev_id 传递给 `device_id` 构造函数参数。
+
+这确保了 FSDP 实例的计算设备是目标设备。对于选项 1 和 3，FSDP 初始化始终在 GPU 上进行。对于选项 2，FSDP 初始化在模块的当前设备上进行，该设备可能是 CPU。
+
+如果你正在使用 `sync_module_states=True` 标志，你需要确保模块在 GPU 上，或者使用 `device_id` 参数指定一个 CUDA 设备，以便 FSDP 在 FSDP 构造函数中将模块移动到该设备。这是必要的，因为 `sync_module_states=True` 需要 GPU 通信。
+
+FSDP 还负责将传递给 forward 方法的输入张量移动到 GPU 计算设备，因此你不需要手动将它们从 CPU 移动。
+
+对于 `use_orig_params=True`，与 `ShardingStrategy.FULL_SHARD` 不同，`ShardingStrategy.SHARD_GRAD_OP` 在前向传播后会暴露未分片的参数，而不是分片后的参数。如果你想检查梯度，可以使用 `summon_full_params` 方法并配合 `with_grads=True`。
+
+在 `limit_all_gathers=True` 的情况下，你可能会在 FSDP 前向传播之前看到一个间隙，此时 CPU 线程没有发出任何内核。这是有意为之，表明速率限制器正在起作用。以这种方式同步 CPU 线程可以防止为后续的全局聚合过度分配内存，并且它实际上不会延迟 GPU 内核的执行。
+
+由于自动微分相关的原因，FSDP 在前向和反向计算期间会用 `torch.Tensor` 视图替换被管理模块的参数。如果你的模块的 forward 依赖于保存的参数引用，而不是在每次迭代时重新获取引用，那么它将看不到 FSDP 新创建的视图，并且自动微分将无法正常工作。
+
+最后，当使用 `sharding_strategy=ShardingStrategy.HYBRID_SHARD` 且分片进程组在节点内、复制进程组在节点间时，对于某些集群设置，设置 `NCCL_CROSS_NIC=1` 可以帮助缩短复制进程组的 all-reduce 时间。
+
+使用 FSDP 时有几个限制需要注意：
+
+当使用 CPU 卸载时，FSDP 目前不支持在 `no_sync()` 之外进行梯度累积。这是因为 FSDP 使用新规约的梯度，而不是与任何现有的梯度进行累积，这可能导致不正确的结果。
+
+FSDP 不支持运行包含在 FSDP 实例中的子模块的前向传播。这是因为子模块的参数将被分片，但子模块本身不是 FSDP 实例，因此它的前向传播将无法适当地 all-gather 完整的参数。
+
+由于其注册反向钩子的方式，FSDP 不支持二阶反向传播（double backwards）。
+
+FSDP 在冻结参数时有一些限制。对于 `use_orig_params=False`，每个 FSDP 实例必须管理全部被冻结或全部未被冻结的参数。对于 `use_orig_params=True`，FSDP 支持混合冻结和未冻结的参数，但建议不要这样做，以防止梯度内存使用量高于预期。
+
+从 PyTorch 1.12 开始，FSDP 提供了对共享参数的有限支持。如果你的用例需要增强的共享参数支持，请在此 issue 中发帖。
+
+你应该避免在不使用 `summon_full_params` 上下文的情况下在前向和反向之间修改参数，因为修改可能不会持久化。
+
+module (nn.Module) – 这是要用 FSDP 包装的模块。
+
+process_group (可选[Union[ProcessGroup, Tuple[ProcessGroup, ProcessGroup]]]) – 这是用于对模型进行分片的进程组，因此也是用于 FSDP 的 all-gather 和 reduce-scatter 集合通信的进程组。如果为 None，则 FSDP 使用默认进程组。对于混合分片策略（如 `ShardingStrategy.HYBRID_SHARD`），用户可以传入一个进程组元组，分别表示用于分片和复制的组。如果为 None，则 FSDP 为用户构造进程组，以在节点内进行分片并在节点间进行复制。(默认: None)
+
+sharding_strategy (可选[ShardingStrategy]) – 这用于配置分片策略，该策略可以在节省内存和通信开销之间进行权衡。详情请参见 `ShardingStrategy`。(默认: FULL_SHARD)
+
+cpu_offload (可选[CPUOffload]) – 这用于配置 CPU 卸载。如果设置为 None，则不会发生 CPU 卸载。详情请参见 `CPUOffload`。(默认: None)
+
+auto_wrap_policy (可选[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy, CustomPolicy]]) – 这指定了将 FSDP 应用于 `module` 的子模块的策略，这对于通信和计算的重叠是必需的，因此会影响性能。如果为 None，则 FSDP 仅应用于 `module`，用户应自行手动将 FSDP 应用于父模块（自底向上进行）。为了方便起见，它直接接受 `ModuleWrapPolicy`，允许用户指定要包装的模块类（例如 transformer 块）。否则，它应该是一个接受三个参数 `module: nn.Module`、`recurse: bool` 和 `nonwrapped_numel: int` 的可调用对象，并且如果 `recurse=False`，应返回一个布尔值指定传入的模块是否应应用 FSDP，或者如果 `recurse=True`，指定是否应继续遍历到该模块的子树中。用户可以向该可调用对象添加额外的参数。`torch.distributed.fsdp.wrap.py` 中的 `size_based_auto_wrap_policy` 给出了一个示例可调用对象，如果模块子树中的参数超过 1 亿（100M）个元素，则将 FSDP 应用于该模块。我们建议在应用 FSDP 后打印模型并根据需要进行调整。 示例: >>> def custom_auto_wrap_policy( >>> module: nn.Module, >>> recurse: bool, >>> nonwrapped_numel: int, >>> # Additional custom arguments >>> min_num_params: int = int(1e8), >>> ) -> bool: >>> return nonwrapped_numel >= min_num_params >>> # Configure a custom `min_num_params` >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
+
+这指定了将 FSDP 应用于 `module` 的子模块的策略，这对于通信和计算的重叠是必需的，因此会影响性能。如果为 None，则 FSDP 仅应用于 `module`，用户应自行手动将 FSDP 应用于父模块（自底向上进行）。为了方便起见，它直接接受 `ModuleWrapPolicy`，允许用户指定要包装的模块类（例如 transformer 块）。否则，它应该是一个接受三个参数 `module: nn.Module`、`recurse: bool` 和 `nonwrapped_numel: int` 的可调用对象，并且如果 `recurse=False`，应返回一个布尔值指定传入的模块是否应应用 FSDP，或者如果 `recurse=True`，指定是否应继续遍历到该模块的子树中。用户可以向该可调用对象添加额外的参数。`torch.distributed.fsdp.wrap.py` 中的 `size_based_auto_wrap_policy` 给出了一个示例可调用对象，如果模块子树中的参数超过 1 亿（100M）个元素，则将 FSDP 应用于该模块。我们建议在应用 FSDP 后打印模型并根据需要进行调整。
+
+backward_prefetch (可选[BackwardPrefetch]) – 这用于配置 all-gather 的显式反向预取。如果为 None，则 FSDP 不进行反向预取，并且在反向传播中没有通信和计算的重叠。详情请参见 `BackwardPrefetch`。(默认: BACKWARD_PRE)
+
+mixed_precision (可选[MixedPrecision]) – 这用于配置 FSDP 的原生混合精度。如果设置为 None，则不使用混合精度。否则，可以设置参数、缓冲区和梯度规约的数据类型。详情请参见 `MixedPrecision`。(默认: None)
+
+ignored_modules (可选[Iterable[torch.nn.Module]]) – 其自身参数以及子模块的参数和缓冲区将被此实例忽略的模块。`ignored_modules` 中直接包含的模块都不应该是 `FullyShardedDataParallel` 实例，并且如果已经是构造好的 `FullyShardedDataParallel` 实例的任何子模块嵌套在此实例下，也不会被忽略。当使用 `auto_wrap_policy` 或者参数的分片不受 FSDP 管理时，此参数可用于在模块粒度上避免对特定参数进行分片。(默认: None)
+
+param_init_fn (可选[Callable[[nn.Module], None]]) – 一个 `Callable[torch.nn.Module] -> None`，指定了当前位于 meta 设备上的模块应如何初始化到实际设备上。从 v1.12 开始，FSDP 通过 `is_meta` 检测在 meta 设备上具有参数或缓冲区的模块，并在指定时应用 `param_init_fn`，否则调用 `nn.Module.reset_parameters()`。对于这两种情况，实现应该只初始化该模块的参数/缓冲区，而不是其子模块的参数/缓冲区。这是为了避免重复初始化。此外，FSDP 还通过 torchdistX 的 (pytorch/torchdistX) `deferred_init()` API 支持延迟初始化，其中延迟的模块通过调用 `param_init_fn`（如果指定）或 torchdistX 的默认 `materialize_module()` 进行初始化。如果指定了 `param_init_fn`，则它将应用于所有 meta 设备模块，这意味着它可能应该根据模块类型进行分支处理。FSDP 在参数展平和分片之前调用初始化函数。 示例: >>> module = MyModule(device="meta") >>> def my_init_fn(module: nn.Module): >>> # E.g. initialize depending on the module type >>> ... >>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy) >>> print(next(fsdp_model.parameters()).device) # current CUDA device >>> # With torchdistX >>> module = deferred_init.deferred_init(MyModule, device="cuda") >>> # Will initialize via deferred_init.materialize_module(). >>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
+
+一个 `Callable[torch.nn.Module] -> None`，指定了当前位于 meta 设备上的模块应如何初始化到实际设备上。从 v1.12 开始，FSDP 通过 `is_meta` 检测在 meta 设备上具有参数或缓冲区的模块，并在指定时应用 `param_init_fn`，否则调用 `nn.Module.reset_parameters()`。对于这两种情况，实现应该只初始化该模块的参数/缓冲区，而不是其子模块的参数/缓冲区。这是为了避免重复初始化。此外，FSDP 还通过 torchdistX 的 (pytorch/torchdistX) `deferred_init()` API 支持延迟初始化，其中延迟的模块通过调用 `param_init_fn`（如果指定）或 torchdistX 的默认 `materialize_module()` 进行初始化。如果指定了 `param_init_fn`，则它将应用于所有 meta 设备模块，这意味着它可能应该根据模块类型进行分支处理。FSDP 在参数展平和分片之前调用初始化函数。
+
+device_id (可选[Union[int, torch.device]]) – 一个 int 或 `torch.device`，指示进行 FSDP 初始化（包括模块初始化（如果需要）和参数分片）所在的 CUDA 设备。如果模块在 CPU 上，应指定此项以提高初始化速度。如果设置了默认 CUDA 设备（例如通过 `torch.cuda.set_device`），则用户可以将 `torch.cuda.current_device` 传递给此参数。(默认: None)
+
+sync_module_states (bool) – 如果为 True，则每个 FSDP 模块将从 rank 0 广播模块参数和缓冲区，以确保它们在各 rank 之间复制（这会给此构造函数增加通信开销）。这有助于通过 `load_state_dict` 以节省内存的方式加载 `state_dict` 检查点。有关此示例，请参见 `FullStateDictConfig`。(默认: False)
+
+forward_prefetch (bool) – 如果为 True，则 FSDP 会在当前前向计算之前显式预取下一次前向传播的 all-gather。这仅对 CPU 密集型工作负载有用，在这种情况下，提前发出下一次 all-gather 可能会改善重叠。这仅适用于静态图模型，因为预取遵循第一次迭代的执行顺序。(默认: False)
+
+limit_all_gathers (bool) – 如果为 True，则 FSDP 会显式同步 CPU 线程，以确保仅由两个连续的 FSDP 实例（当前正在运行计算的实例和正在预取 all-gather 的下一个实例）占用 GPU 内存。如果为 False，则 FSDP 允许 CPU 线程发出 all-gather，而无需任何额外的同步。(默认: True) 我们通常将此功能称为“速率限制器”。仅对于内存压力低且受 CPU 限制的特定工作负载，才应将此标志设置为 False，在这种情况下，CPU 线程可以激进地发出所有内核，而无需担心 GPU 内存使用量。
+
+use_orig_params (bool) – 将此设置为 True 会让 FSDP 使用模块的原始参数。FSDP 通过 `nn.Module.named_parameters()` 将这些原始参数暴露给用户，而不是 FSDP 内部的 `FlatParameter`。这意味着优化器步骤在原始参数上运行，从而允许按原始参数设置超参数。FSDP 保留原始参数变量，并在未分片和分片形式之间操作它们的数据，它们分别始终是底层未分片或分片 `FlatParameter` 的视图。按照当前的算法，分片形式始终是一维的，丢失了原始的张量结构。对于给定的 rank，一个原始参数可能包含其全部、部分或不含任何数据。在不包含数据的情况下，它的数据将类似于一个大小为 0 的空张量。用户不应编写依赖于给定原始参数在其分片形式中存在什么数据的程序。使用 `torch.compile()` 需要将其设置为 True。将此设置为 False 会通过 `nn.Module.named_parameters()` 将 FSDP 内部的 `FlatParameter` 暴露给用户。(默认: False)
+
+ignored_states (可选[Iterable[torch.nn.Parameter]], 可选[Iterable[torch.nn.Module]]) – 不由此 FSDP 实例管理的被忽略参数或模块，这意味着这些参数不会被分片，并且它们的梯度不会在各 rank 之间进行规约。此参数与现有的 `ignored_modules` 参数统一，我们可能很快弃用 `ignored_modules`。为了向后兼容，我们保留了 `ignored_states` 和 `ignored_modules`，但 FSDP 只允许将其中之一指定为非 None。
+
+device_mesh (可选[DeviceMesh]) – `DeviceMesh` 可用作 `process_group` 的替代方法。当传递 `device_mesh` 时，FSDP 将使用底层进程组进行 all-gather 和 reduce-scatter 集合通信。因此，这两个参数必须是互斥的。对于混合分片策略（如 `ShardingStrategy.HYBRID_SHARD`），用户可以传入一个二维 `DeviceMesh` 而不是进程组元组。对于二维 FSDP + TP，用户需要传入 `device_mesh` 而不是 `process_group`。有关 `DeviceMesh` 的更多信息，请访问：https://pytorch.org/tutorials/recipes/distributed_device_mesh.html
+
+将 `fn` 递归地应用于每个子模块（由 `.children()` 返回）以及自身。
+
+典型用途包括初始化模型的参数（另见 `torch.nn.init`）。
+
+与 `torch.nn.Module.apply` 相比，此版本在应用 `fn` 之前会额外收集完整的参数。不应在另一个 `summon_full_params` 上下文中调用此方法。
+
+fn (Module -> None) – 应用于每个子模块的函数
+
+检查此实例是否为根 FSDP 模块。
+
+裁剪所有参数的梯度范数。
+
+该范数是将在所有参数的梯度（视为单个向量）上计算的，并且梯度会被原地修改。
+
+max_norm (float or int) – 梯度的最大范数
+
+norm_type (float or int) – 使用的 p 范数的类型。可以为 'inf' 表示无穷范数。
+
+参数的总范数（视为单个向量）。
+
+如果每个 FSDP 实例都使用 `NO_SHARD`，这意味着没有梯度在各 rank 之间进行分片，那么你可以直接使用 `torch.nn.utils.clip_grad_norm_()`。
+
+如果至少有一些 FSDP 实例使用了分片策略（即非 `NO_SHARD` 的策略），那么你应该使用此方法而不是 `torch.nn.utils.clip_grad_norm_()`，因为此方法处理了梯度在各 rank 之间分片的事实。
+
+返回的总范数将具有根据 PyTorch 的类型提升语义定义的所有参数/梯度中“最大”的数据类型。例如，如果所有参数/梯度都使用低精度数据类型，则返回范数的数据类型将是该低精度数据类型，但如果存在至少一个使用 FP32 的参数/梯度，则返回范数的数据类型将为 FP32。
+
+由于使用集合通信，因此需要在所有 rank 上调用此方法。
+
+展平分片的优化器状态字典（state-dict）。
+
+该 API 类似于 `shard_full_optim_state_dict()`。唯一的区别在于，输入的 `sharded_optim_state_dict` 应从 `sharded_optim_state_dict()` 返回。因此，每个 rank 上将会有 all-gather 调用来收集 `ShardedTensor`。
+
+sharded_optim_state_dict (Dict[str, Any]) – 对应于未展平参数并保存分片优化器状态的优化器状态字典。
+
+model (torch.nn.Module) – 参考 `shard_full_optim_state_dict()`。
+
+optim (torch.optim.Optimizer) – 针对 `model` 参数的优化器。
+
+参考 `shard_full_optim_state_dict()`。
+
+运行被包装模块的前向传播，插入 FSDP 特有的前向前和前向后分片逻辑。
+
+返回所有嵌套的 FSDP 实例。
+
+这可能包含模块本身，并且如果 `root_only=True`，则仅包含 FSDP 根模块。
+
+module (torch.nn.Module) – 根模块，可能是也可能不是 FSDP 模块。
+
+root_only (bool) – 是否仅返回 FSDP 根模块。(默认: False)
+
+嵌套在输入模块中的 FSDP 模块。
+
+List[FullyShardedDataParallel]
+
+返回完整的优化器状态字典。
+
+在 rank 0 上合并完整的优化器状态，并按照 `torch.optim.Optimizer.state_dict()` 的约定将其作为字典返回（即带有键 "state" 和 "param_groups"）。包含在 `model` 中的 FSDP 模块中已展平的参数将被映射回其未展平的参数。
+
+由于使用集合通信，因此需要在所有 rank 上调用此方法。但是，如果 `rank0_only=True`，则仅在 rank 0 上填充状态字典，所有其他 rank 返回一个空字典。
+
+与 `torch.optim.Optimizer.state_dict()` 不同，此方法使用完整的参数名称作为键，而不是参数 ID。
+
+与 `torch.optim.Optimizer.state_dict()` 类似，优化器状态字典中包含的张量没有被克隆，因此可能会出现别名问题。作为最佳实践，考虑立即保存返回的优化器状态字典，例如使用 `torch.save()`。
+
+model (torch.nn.Module) – 根模块（可能是也可能不是 `FullyShardedDataParallel` 实例），其参数被传递给优化器 `optim`。
+
+optim (torch.optim.Optimizer) – 针对 `model` 参数的优化器。
+
+optim_input (可选[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递给优化器 `optim` 的输入，表示参数组列表或可迭代的参数；如果为 None，则此方法假定输入为 `model.parameters()`。此参数已弃用，无需再传入。(默认: None)
+
+rank0_only (bool) – 如果为 True，则仅在 rank 0 上保存填充的字典；如果为 False，则在所有 rank 上保存。(默认: True)
+
+group (dist.ProcessGroup) – 模型的进程组，如果使用默认进程组则为 None。(默认: None)
+
+一个包含 `model` 原始未展平参数的优化器状态的字典，并且遵循 `torch.optim.Optimizer.state_dict()` 的约定包含键 “state” 和 “param_groups”。如果 `rank0_only=True`，则非零 rank 返回一个空字典。
+
+获取根植于 `module` 的 FSDP 模块的 `state_dict_type` 及相应的配置。
+
+目标模块不必是 FSDP 模块。
+
+一个 `StateDictSettings`，包含当前设置的 `state_dict_type` 和 `state_dict` / `optim_state_dict` 配置。
+
+如果不同 FSDP 子模块的 `StateDictSettings` 不同，则抛出 `AssertionError`。
+
+返回被包装的模块。
+
+返回模块缓冲区的迭代器，生成缓冲区的名称和缓冲区本身。
+
+在 `summon_full_params()` 上下文管理器内，拦截缓冲区名称并移除所有出现的 FSDP 特定展平缓冲区前缀。
+
+Iterator[tuple[str, torch.Tensor]]
+
+返回模块参数的迭代器，生成参数的名称和参数本身。
+
+在 `summon_full_params()` 上下文管理器内，拦截参数名称并移除所有出现的 FSDP 特定展平参数前缀。
+
+Iterator[tuple[str, torch.nn.parameter.Parameter]]
+
+禁用 FSDP 实例之间的梯度同步。
+
+在此上下文中，梯度将在模块变量中累积，稍后在退出上下文后的第一次前向-反向传播中进行同步。这仅应在根 FSDP 实例上使用，并将递归应用于所有子 FSDP 实例。
+
+这可能会导致更高的内存使用量，因为 FSDP 会累积完整的模型梯度（而不是梯度分片），直到最终同步。
+
+当与 CPU 卸载一起使用时，在上下文管理器内梯度不会被卸载到 CPU。相反，它们只会在最终同步之后被卸载。
+
+转换对应于分片模型的优化器的 state-dict。
+
+给定的 state-dict 可以转换为三种类型之一：1) 完整的优化器 state_dict，2) 分片的优化器 state_dict，3) 本地优化器 state_dict。
+
+对于完整的优化器 state_dict，所有状态都是未展平且未分片的。可以通过 `state_dict_type()` 指定仅 rank0 和仅 CPU，以避免 OOM。
+
+对于分片的优化器 state_dict，所有状态都是未展平但已分片的。可以通过 `state_dict_type()` 指定仅 CPU 以进一步节省内存。
+
+对于本地 state_dict，不会执行任何转换。但状态将从 `nn.Tensor` 转换为 `ShardedTensor` 以表示其分片性质（目前尚不支持）。
+
+model (torch.nn.Module) – 根模块（可能是也可能不是 `FullyShardedDataParallel` 实例），其参数被传递给优化器 `optim`。
+
+optim (torch.optim.Optimizer) – 针对 `model` 参数的优化器。
+
+optim_state_dict (Dict[str, Any]) – 要转换的目标优化器 state_dict。如果值为 None，将使用 `optim.state_dict()`。(默认: None)
+
+group (dist.ProcessGroup) – 模型的进程组，参数在该组上进行分片；如果使用默认进程组，则为 None。(默认: None)
+
+一个包含 `model` 的优化器状态的字典。优化器状态的分片基于 `state_dict_type`。
+
+转换优化器状态字典，以便将其加载到与 FSDP 模型关联的优化器中。
+
+给定一个通过 `optim_state_dict()` 转换的 `optim_state_dict`，它将被转换为可以加载到 `optim`（即 `model` 的优化器）中的展平优化器状态字典。`model` 必须通过 `FullyShardedDataParallel` 进行分片。
+
+model (torch.nn.Module) – 根模块（可能是也可能不是 `FullyShardedDataParallel` 实例），其参数被传递给优化器 `optim`。
+
+optim (torch.optim.Optimizer) – 针对 `model` 参数的优化器。
+
+optim_state_dict (Dict[str, Any]) – 要加载的优化器状态。
+
+is_named_optimizer (bool) – 该优化器是否为 `NamedOptimizer` 或 `KeyedOptimizer`。仅当 `optim` 是 TorchRec 的 `KeyedOptimizer` 或 `torch.distributed` 的 `NamedOptimizer` 时才设置为 True。
+
+load_directly (bool) – 如果设置为 True，此 API 将在返回结果之前调用 `optim.load_state_dict(result)`。否则，用户需自行负责调用 `optim.load_state_dict()` (默认: False)
+
+group (dist.ProcessGroup) – 模型的进程组，参数在该组上进行分片；如果使用默认进程组，则为 None。(默认: None)
+
+注册通信钩子。
+
+这是一个增强功能，为用户提供了灵活的钩子，使他们可以指定 FSDP 如何在多个工作进程之间聚合梯度。此钩子可用于实现多种算法，如 GossipGrad 和梯度压缩，这些算法涉及在使用 `FullyShardedDataParallel` 进行训练时用于参数同步的不同通信策略。
+
+FSDP 通信钩子应在运行初始前向传播之前注册，并且仅能注册一次。
+
+state (object) – 传递给钩子以在训练过程中维护任何状态信息。示例包括梯度压缩中的误差反馈、GossipGrad 中下一个要与之通信的对等节点等。它由每个工作进程本地存储，并由该工作进程上的所有梯度张量共享。
+
+传递给钩子以在训练过程中维护任何状态信息。示例包括梯度压缩中的误差反馈、GossipGrad 中下一个要与之通信的对等节点等。它由每个工作进程本地存储，并由该工作进程上的所有梯度张量共享。
+
+hook (Callable) – 可调用对象，具有以下签名之一：1) `hook: Callable[torch.Tensor] -> None`：此函数接收一个 Python 张量，它表示相对于此 FSDP 单元所包装的模型对应的所有变量（未被其他 FSDP 子单元包装的）的完整的、展平的、未分片的梯度。然后它执行所有必要的处理并返回 None；2) `hook: Callable[torch.Tensor, torch.Tensor] -> None`：此函数接收两个 Python 张量，第一个表示相对于此 FSDP 单元所包装的模型对应的所有变量（未被其他 FSDP 子单元包装的）的完整的、展平的、未分片的梯度。后者表示一个预先设定大小的张量，用于在规约后存储一块分片的梯度。在这两种情况下，可调用对象执行所有必要的处理并返回 None。签名为 1 的可调用对象预期用于处理 `NO_SHARD` 情况下的梯度通信。签名为 2 的可调用对象预期用于处理分片情况下的梯度通信。
+
+重新键控（Re-keys）优化器状态字典 `optim_state_dict`，以使用 `optim_state_key_type` 指定的键类型。
+
+这可用于在具有 FSDP 实例的模型和不具有 FSDP 实例的模型之间的优化器状态字典之间实现兼容性。
+
+要将 FSDP 完整优化器状态字典（即来自 `full_optim_state_dict()`）重新键控为使用参数 ID，并使其可加载到非包装的模型中：
+
+要将来自非包装模型的普通优化器状态字典重新键控为可加载到包装的模型中：
+
+使用 `optim_state_key_type` 指定的参数键重新键控的优化器状态字典。
+
+将完整的优化器状态字典从 rank 0 分发到所有其他 rank。
+
+在每个 rank 上返回分片的优化器状态字典。返回值与 `shard_full_optim_state_dict()` 相同，在 rank 0 上，第一个参数应该是 `full_optim_state_dict()` 的返回值。
+
+`shard_full_optim_state_dict()` 和 `scatter_full_optim_state_dict()` 都可用于获取要加载的分片优化器状态字典。假设完整的优化器状态字典位于 CPU 内存中，前者要求每个 rank 在 CPU 内存中都有完整的字典，每个 rank 独立地对字典进行分片而无需任何通信；而后者仅要求 rank 0 在 CPU 内存中拥有完整字典，rank 0 将每个分片移动到 GPU 内存（用于 NCCL）并将其适当地通信给各个 rank。因此，前者的总体 CPU 内存成本较高，而后者的通信成本较高。
+
+full_optim_state_dict (可选[Dict[str, Any]]) – 对应于未展平参数并保存完整非分片优化器状态的优化器状态字典（如果在 rank 0 上）；在非零 rank 上忽略此参数。
+
+model (torch.nn.Module) – 根模块（可能是也可能不是 `FullyShardedDataParallel` 实例），其参数对应于 `full_optim_state_dict` 中的优化器状态。
+
+optim_input (可选[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递给优化器的输入，表示参数组列表或可迭代的参数；如果为 None，则此方法假定输入为 `model.parameters()`。此参数已弃用，无需再传入。(默认: None)
+
+optim (可选[torch.optim.Optimizer]) – 将加载由此方法返回的状态字典的优化器。这是比 `optim_input` 更推荐使用的参数。(默认: None)
+
+group (dist.ProcessGroup) – 模型的进程组，如果使用默认进程组则为 None。(默认: None)
+
+现在重新映射为展平参数（而不是未展平参数）并且仅包含此 rank 部分优化器状态的完整优化器状态字典。
+
+设置目标模块的所有后代 FSDP 模块的 `state_dict_type`。
+
+还可以为模型和优化器的状态字典进行（可选的）配置。目标模块不必是 FSDP 模块。如果目标模块是 FSDP 模块，其 `state_dict_type` 也将被更改。
+
+此 API 应仅用于顶层（根）模块。
+
+此 API 使用户能够在根 FSDP 模块被另一个 `nn.Module` 包装的情况下，透明地使用传统的 `state_dict` API 来获取模型检查点。例如，以下操作将确保在所有非 FSDP 实例上调用 `state_dict`，同时为 FSDP 调度到 `sharded_state_dict` 实现：
+
+module (torch.nn.Module) – 根模块。
+
+state_dict_type (StateDictType) – 要设置的期望的 `state_dict_type`。
+
+state_dict_config (可选[StateDictConfig]) – 目标 `state_dict_type` 的配置。
+
+optim_state_dict_config (可选[OptimStateDictConfig]) – 优化器状态字典的配置。
+
+一个 `StateDictSettings`，包含模块先前的 `state_dict` 类型和配置。
+
+分片完整的优化器状态字典。
+
+将 `full_optim_state_dict` 中的状态重新映射为展平参数（而不是未展平参数），并仅限制为此 rank 部分的优化器状态。第一个参数应该是 `full_optim_state_dict()` 的返回值。
+
+`shard_full_optim_state_dict()` 和 `scatter_full_optim_state_dict()` 都可用于获取要加载的分片优化器状态字典。假设完整的优化器状态字典位于 CPU 内存中，前者要求每个 rank 在 CPU 内存中都有完整的字典，每个 rank 独立地对字典进行分片而无需任何通信；而后者仅要求 rank 0 在 CPU 内存中拥有完整字典，rank 0 将每个分片移动到 GPU 内存（用于 NCCL）并将其适当地通信给各个 rank。因此，前者的总体 CPU 内存成本较高，而后者的通信成本较高。
+
+full_optim_state_dict (Dict[str, Any]) – 对应于未展平参数并保存完整非分片优化器状态的优化器状态字典。
+
+model (torch.nn.Module) – 根模块（可能是也可能不是 `FullyShardedDataParallel` 实例），其参数对应于 `full_optim_state_dict` 中的优化器状态。
+
+optim_input (可选[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递给优化器的输入，表示参数组列表或可迭代的参数；如果为 None，则此方法假定输入为 `model.parameters()`。此参数已弃用，无需再传入。(默认: None)
+
+optim (可选[torch.optim.Optimizer]) – 将加载由此方法返回的状态字典的优化器。这是比 `optim_input` 更推荐使用的参数。(默认: None)
+
+现在重新映射为展平参数（而不是未展平参数）并且仅包含此 rank 部分优化器状态的完整优化器状态字典。
+
+返回分片形式的优化器状态字典。
+
+该 API 类似于 `full_optim_state_dict()`，但此 API 将所有非零维度的状态分块为 `ShardedTensor` 以节省内存。仅当模型状态字典是通过 `with state_dict_type(SHARDED_STATE_DICT):` 上下文管理器派生时，才应使用此 API。
+
+有关详细用法，请参考 `full_optim_state_dict()`。
+
+返回的状态字典包含 `ShardedTensor`，不能直接被常规的 `optim.load_state_dict` 使用。
+
+设置目标模块的所有后代 FSDP 模块的 `state_dict_type`。
+
+此上下文管理器具有与 `set_state_dict_type()` 相同的功能。有关详细信息，请阅读 `set_state_dict_type()` 的文档。
+
+module (torch.nn.Module) – 根模块。
+
+state_dict_type (StateDictType) – 要设置的期望的 `state_dict_type`。
+
+state_dict_config (可选[StateDictConfig]) – 针对 `state_dict_type` 的模型状态字典配置。
+
+optim_state_dict_config (可选[OptimStateDictConfig]) – 针对 `state_dict_type` 的优化器状态字典配置。
+
+使用此上下文管理器为 FSDP 实例暴露完整参数。
+
+在前向/反向传播之后，这对模型获取参数以进行额外处理或检查可能很有用。它可以接受非 FSDP 模块，并将根据 `recurse` 参数为所有包含的 FSDP 模块及其子模块调用全部参数。
+
+这可以在内部 FSDP 上使用。
+
+这不能在前向或反向传播中使用。前向和反向也不能从此上下文中开始。
+
+上下文管理器退出后，参数将恢复为其本地分片，存储行为与前向传播相同。
+
+可以修改完整参数，但只有与本地参数分片相对应的部分会在上下文管理器退出后保留（除非 `writeback=False`，在这种情况下更改将被丢弃）。在 FSDP 不对参数进行分片的情况下（目前仅在 `world_size == 1` 或 `NO_SHARD` 配置时），无论 `writeback` 如何，修改都会保留。
+
+此方法适用于本身不是 FSDP 但可能包含多个独立 FSDP 单元的模块。在这种情况下，给定的参数将应用于所有包含的 FSDP 单元。
+
+请注意，目前不支持 `rank0_only=True` 结合 `writeback=True`，并且会引发错误。这是因为在此上下文中，模型参数的形状在各 rank 之间会不同，向它们写入数据可能会导致退出上下文时各 rank 之间的不一致。
+
+请注意，`offload_to_cpu` 和 `rank0_only=False` 将导致位于同一台机器上的 GPU 将完整参数冗余复制到 CPU 内存中，这可能会导致 CPU OOM 的风险。建议使用 `offload_to_cpu` 并结合 `rank0_only=True`。
+
+recurse (bool, 可选) – 递归调用嵌套 FSDP 实例的所有参数 (默认: True)。
+
+writeback (bool, 可选) – 如果为 False，则在上下文管理器退出后放弃对参数的修改；禁用此选项可能会稍微提高效率 (默认: True)
+
+rank0_only (bool, 可选) – 如果为 True，则仅在全局 rank 0 上实例化完整参数。这意味着在此上下文中，只有 rank 0 将拥有完整参数，而其他 rank 将拥有分片参数。请注意，不支持将 `rank0_only=True` 与 `writeback=True` 一起使用，因为在此上下文中，模型参数的形状在各 rank 之间会不同，向它们写入数据可能会导致退出上下文时各 rank 之间的不一致。
+
+offload_to_cpu (bool, 可选) – 如果为 True，则将完整参数卸载到 CPU。请注意，目前仅当参数被分片时（仅在 `world_size = 1` 或 `NO_SHARD` 配置下不是这种情况）才会发生此卸载。建议使用 `offload_to_cpu` 并结合 `rank0_only=True`，以避免将模型参数的冗余副本卸载到相同的 CPU 内存中。
+
+with_grads (bool, 可选) – 如果为 True，梯度也会与参数一起取消分片。目前，仅当向 FSDP 构造函数传递 `use_orig_params=True` 且向此方法传递 `offload_to_cpu=False` 时，才支持此选项。(默认: False)
+
+这配置了显式的反向预取，它通过在反向传播中实现通信和计算重叠来提高吞吐量，代价是内存使用量略微增加。
+
+BACKWARD_PRE: 这启用了最大程度的重叠，但内存使用量增加最多。这在当前一组参数的梯度计算之前预取下一组参数。这使下一次 all-gather 和当前的梯度计算重叠，并且在峰值时，它会在内存中保存当前这组参数、下一组参数
+
+## 分布式优化器#
+
+**URL:** https://pytorch.org/docs/stable/distributed.optim.html
+
+**目录：**
+- 分布式优化器#
+
+创建于：2021年3月1日 | 最后更新于：2025年6月16日
+
+当使用 CUDA 张量时，目前不支持分布式优化器
+
+torch.distributed.optim 公开了 `DistributedOptimizer`，它接收一个远程参数列表（`RRef`），并在参数所在的工作节点上本地运行优化器。分布式优化器可以使用任何本地优化器基类在每个工作节点上应用梯度。
+
+`DistributedOptimizer` 接收分布在各个工作节点上的参数的远程引用，并在本地为每个参数应用给定的优化器。
+
+该类使用 `get_gradients()` 来检索特定参数的梯度。
+
+对 `step()` 的并发调用（无论是来自同一客户端还是不同客户端）都将在每个工作节点上被串行化——因为每个工作节点的优化器一次只能处理一组梯度。但是，不能保证完整的“前向-反向-优化器”序列会一次只针对一个客户端执行。这意味着，所应用的梯度可能不对应于在给定工作节点上执行的最新前向传播。此外，跨工作节点也没有保证的顺序。
+
+`DistributedOptimizer` 默认启用 TorchScript 创建本地优化器，以便在多线程训练（例如分布式模型并行）的情况下，优化器的更新不会被 Python 全局解释器锁（GIL）阻塞。此功能目前已为大多数优化器启用。你也可以按照 PyTorch 教程中的指南，为你自己的自定义优化器启用 TorchScript 支持。
+
+optimizer_class (optim.Optimizer) – 在每个工作节点上实例化的优化器类。
+
+params_rref (list[RRef]) – 要优化的本地或远程参数的 RRef 列表。
+
+args – 传递给每个工作节点上优化器构造函数的参数。
+
+kwargs – 传递给每个工作节点上优化器构造函数的参数。
+
+执行单次优化步骤。
+
+这将在每个包含待优化参数的工作节点上调用 `torch.optim.Optimizer.step()`，并阻塞直到所有工作节点返回。提供的 `context_id` 将用于检索包含应应用于参数的梯度的相应上下文。
+
+context_id – 我们应为其运行优化器步骤的 autograd 上下文 ID。
+
+包装任意的 `torch.optim.Optimizer` 并运行 post-local SGD，此优化器在每一步运行本地优化器。在预热阶段之后，它会在应用本地优化器后定期对参数进行平均。
+
+optim (Optimizer) – 本地优化器。
+
+averager (ModelAverager) – 用于运行 post-localSGD 算法的模型平均器实例。
+
+这与 `torch.optim.Optimizer` 的 `load_state_dict()` 相同，但还会将模型平均器的步数值恢复为提供的 `state_dict` 中保存的值。
+
+如果 `state_dict` 中没有 "step" 条目，它将发出警告并将模型平均器的步数初始化为 0。
+
+这与 `torch.optim.Optimizer` 的 `state_dict()` 相同，但添加了一个额外的条目，用于将模型平均器的步数记录到检查点中，以确保重新加载时不会再次导致不必要的预热。
+
+执行单次优化步骤（参数更新）。
+
+包装任意的 `optim.Optimizer` 并将其状态分片到组中的各个 rank 上。
+
+共享机制按照 ZeRO 的描述进行。
+
+每个 rank 中的本地优化器实例仅负责更新大约 1 / world_size 的参数，因此只需保留 1 / world_size 的优化器状态。在本地更新参数之后，每个 rank 会将其参数广播给所有其他节点，以保持所有模型副本处于相同状态。`ZeroRedundancyOptimizer` 可以与 `torch.nn.parallel.DistributedDataParallel` 结合使用，以降低每个 rank 的峰值内存消耗。
+
+`ZeroRedundancyOptimizer` 使用排序贪心算法在每个 rank 上打包多个参数。每个参数仅属于单个 rank，并且不会在 rank 之间进行切分。这种划分是任意的，可能不匹配参数的注册或使用顺序。
+
+params (Iterable) – 包含所有参数的 `torch.Tensor` 或 `dict` 的可迭代对象，这些参数将跨 rank 进行分片。
+
+optimizer_class (torch.nn.Optimizer) – 本地优化器的类。
+
+process_group (ProcessGroup, 可选) – `torch.distributed` ProcessGroup（默认值：由 `torch.distributed.init_process_group()` 初始化的 `dist.group.WORLD`）。
+
+parameters_as_bucket_view (bool, 可选) – 如果为 True，参数被打包到桶中以加速通信，并且 `param.data` 字段指向不同偏移量的桶视图；如果为 False，每个单独的参数将单独通信，并且每个 `params.data` 保持不变（默认值：False）。
+
+overlap_with_ddp (bool, 可选) – 如果为 True，`step()` 将与 `DistributedDataParallel` 的梯度同步重叠；这要求（1）`optimizer_class` 参数要么是函数式优化器，要么是其具有函数式等价物的优化器，并且（2）注册一个由 `ddp_zero_hook.py` 中的函数之一构建的 DDP 通信钩子；参数被打包到与 `DistributedDataParallel` 中的桶相匹配的桶中，这意味着 `parameters_as_bucket_view` 参数将被忽略。如果为 False，`step()` 将在反向传播之后独立运行（按正常情况）。（默认值：False）
+
+**defaults – 任何尾随参数，将直接转发给本地优化器。
+
+目前，`ZeroRedundancyOptimizer` 要求传入的所有参数都是相同的稠密类型。
+
+如果传入 `overlap_with_ddp=True` 请注意以下事项：鉴于目前将 `DistributedDataParallel` 与 `ZeroRedundancyOptimizer` 重叠的实现方式，前两到三次训练迭代不会在优化器步骤中执行参数更新，具体分别取决于 `static_graph=False` 还是 `static_graph=True`。这是因为它需要获取有关 `DistributedDataParallel` 使用的梯度分桶策略的信息，而这些信息在 `static_graph=False` 时直到第二次前向传播才会最终确定，或者在 `static_graph=True` 时直到第三次前向传播才会最终确定。为了对此进行调整，一种选择是添加虚拟输入。
+
+`ZeroRedundancyOptimizer` 处于实验阶段，可能会发生更改。
+
+向 `Optimizer` 的 `param_groups` 添加参数组。
+
+这在微调预训练网络时非常有用，因为随着训练的进行，被冻结的层可以变为可训练状态，并添加到 `Optimizer` 中。
+
+param_group (dict) – 指定要优化的参数以及特定于组的优化选项。
+
+此方法处理更新所有分区上的分片，但必须在所有 rank 上调用。如果在部分 rank 上调用，会导致训练挂起，因为通信原语是根据受管参数调用的，并且期望所有 rank 都在同一组参数上参与。
+
+将 `state_dict` 列表（每个 rank 一个）合并到目标 rank 上。
+
+to (int) – 接收优化器状态的 rank（默认值：0）。
+
+RuntimeError – 如果 `overlap_with_ddp=True` 并且在此 `ZeroRedundancyOptimizer` 实例完全初始化之前调用此方法（一旦 `DistributedDataParallel` 梯度桶被重建，此实例即被完全初始化）。
+
+这必须在所有 rank 上调用。
+
+返回默认设备。
+
+返回 ZeRO join hook。
+
+它通过遮蔽优化器步骤中的集合通信来支持在不均匀输入上进行训练。
+
+在调用此钩子之前必须正确设置梯度。
+
+kwargs (dict) – 一个包含在运行时修改 join hook 行为的任何关键字参数的字典；所有共享同一 join 上下文管理器的可加入实例都会为 `kwargs` 转发相同的值。
+
+此 hook 不支持任何关键字参数；即 `kwargs` 未被使用。
+
+返回进程组。
+
+从输入的 `state_dict` 加载与给定 rank 相关的状态，并根据需要更新本地优化器。
+
+state_dict (dict) – 优化器状态；应该是从调用 `state_dict()` 返回的对象。
+
+RuntimeError – 如果 `overlap_with_ddp=True` 并且在此 `ZeroRedundancyOptimizer` 实例完全初始化之前调用此方法（一旦 `DistributedDataParallel` 梯度桶被重建，此实例即被完全初始化）。
+
+返回此 rank 已知的最后一次全局优化器状态。
+
+RuntimeError – 如果 `overlap_with_ddp=True` 并且在此 `ZeroRedundancyOptimizer` 实例完全初始化之前调用此方法（一旦 `DistributedDataParallel` 梯度桶被重建，此实例即被完全初始化）；或者在调用 `consolidate_state_dict()` 之前调用了此方法。
+
+执行单次优化器步骤并跨所有 rank 同步参数。
+
+closure (Callable) – 一个重新评估模型并返回 loss 的闭包；对于大多数优化器是可选的。
+
+可选的 loss，具体取决于底层的本地优化器。
+
+任何额外的参数都将原样传递给基础优化器。
+
+---
+
+## Torch 分布式弹性#
+
+**URL:** https://pytorch.org/docs/stable/distributed.elastic.html
+
+**目录：**
+- Torch 分布式弹性#
+- 快速开始#
+- 文档#
+
+创建于：2025年6月16日 | 最后更新于：2025年7月25日
+
+使分布式 PyTorch 具备容错性和弹性。
+
+---
+
+## 流水线并行#
+
+**URL:** https://pytorch.org/docs/stable/distributed.pipelining.html
+
+**目录：**
+- 流水线并行#
+- 为什么需要流水线并行？#
+- 什么是 torch.distributed.pipelining？#
+- 步骤 1：构建 PipelineStage#
+- 步骤 2：使用 PipelineSchedule 执行#
+- 拆分模型的选项#
+  - 选项 1：手动拆分模型#
+  - 选项 2：自动拆分模型#
+- Hugging Face 示例#
+- 技术深入剖析#
+
+创建于：2025年6月16日 | 最后更新于：2025年8月13日
+
+`torch.distributed.pipelining` 目前处于 alpha 阶段，正在开发中。API 可能会发生变化。它是从 PiPPy 项目迁移而来的。
+
+流水线并行是深度学习的基本并行方式之一。它允许将模型的执行进行分区，使得多个微批次可以并发执行模型代码的不同部分。流水线并行对于以下情况可能是一种有效的技术：
+
+带宽受限的集群
+
+大模型推理
+
+上述场景都有一个共同点，即每个设备上的计算无法掩盖传统并行的通信开销，例如 FSDP 的权重全收集。
+
+虽然在扩展方面很有前景，但流水线通常难以实现，因为除了模型权重之外，它还需要对模型的执行进行分区。对执行进行分区通常需要对模型进行侵入性的代码修改。另一个复杂性来自于在分布式环境中调度微批次时需要考虑数据流依赖性。
+
+`pipelining` 包提供了一个能够自动执行上述操作的工具包，从而允许在通用模型上轻松实现流水线并行。
+
+它由两部分组成：拆分前端和分布式运行时。拆分前端接收你的原始模型代码，将其拆分为“模型分区”，并捕获数据流关系。分布式运行时在不同设备上并行执行流水线阶段，处理微批次拆分、调度、通信和梯度传播等事宜。
+
+总体而言，`pipelining` 包提供以下功能：
+
+基于简单的规范拆分模型代码。
+
+对流水线调度提供丰富的支持，包括 GPipe、1F1B、Interleaved 1F1B 和 Looped BFS，并提供编写自定义调度的基础设施。
+
+对跨主机流水线并行提供一流的支持，因为这通常是使用 PP 的场景（通过较慢的互连）。
+
+与其他 PyTorch 并行技术（如数据并行（DDP，FSDP）或张量并行）的可组合性。TorchTitan 项目演示了在 Llama 模型上的“3D 并行”应用。
+
+在我们可以使用 `PipelineSchedule` 之前，我们需要创建 `PipelineStage` 对象来包装在该阶段中运行的模型部分。`PipelineStage` 负责分配通信缓冲区并创建发送/接收操作以与其对等端通信。它管理中间缓冲区（例如用于尚未被消耗的前向传播输出），并提供为阶段模型运行反向传播的工具。
+
+`PipelineStage` 需要知道阶段模型的输入和输出形状，以便它可以正确分配通信缓冲区。这些形状必须是静态的，例如在运行时，形状不能在不同的步骤之间发生变化。如果运行时的形状与预期的形状不匹配，将引发 `PipeliningShapeError` 异常。当与其他并行机制组合或应用混合精度时，必须考虑到这些技术，以便 `PipelineStage` 知道阶段模块在运行时输出的正确形状（和 dtype）。
+
+用户可以直接构造 `PipelineStage` 实例，方法是传入一个 `nn.Module`，该模块表示应在该阶段上运行的那部分模型。这可能需要修改原始模型代码。请参阅选项 1：手动拆分模型中的示例。
+
+或者，拆分前端可以使用图分区自动将你的模型拆分为一系列 `nn.Module`。此技术要求模型可以通过 `torch.Export` 进行跟踪。生成的 `nn.Module` 与其他并行技术的可组合性是实验性的，可能需要一些变通方法。如果用户无法轻易更改模型代码，使用该前端可能会更具吸引力。有关更多信息，请参阅选项 2：自动拆分模型。
+
+现在我们可以将 `PipelineStage` 附加到流水线调度上，并使用输入数据运行调度。以下是一个 GPipe 示例：
+
+请注意，上述代码需要为每个工作节点启动，因此我们使用启动器服务来启动多个进程：
+
+为了直接构造 `PipelineStage`，用户负责提供一个包含相关 `nn.Parameters` 和 `nn.Buffers` 的单一 `nn.Module` 实例，并定义一个执行与该阶段相关的操作的 `forward()` 方法。例如，Torchtitan 中定义的 Transformer 类的精简版本展示了构建易于分区模型的模式。
+
+以这种方式定义的模型可以通过以下方式轻松地为每个阶段进行配置：首先初始化整个模型（使用 meta-device 以避免 OOM 错误），删除该阶段不需要的层，然后创建一个包装该模型的 `PipelineStage`。例如：
+
+当与其他数据或模型并行技术组合时，如果模型块的输出形状/dtype 会受到影响，也可能需要 `output_args`。
+
+如果你有一个完整的模型，并且不想花时间将其修改为一系列“模型分区”，流水线 API 可以为你提供帮助。这是一个简短的示例：
+
+如果我们打印该模型，我们可以看到多个层级，这使得手动拆分变得困难：
+
+让我们看看流水线 API 是如何工作的：
+
+流水线 API 根据给定的 `split_spec` 拆分你的模型，其中 `SplitPoint.BEGINNING` 表示在 forward 函数中某个子模块执行之前添加一个分割点，类似地，`SplitPoint.END` 表示在其执行之后添加一个分割点。
+
+如果我们打印 `pipe`，我们可以看到：
+
+“模型分区”由子模块（`submod_0`，`submod_1`）表示，每个子模块都是用原始模型操作、权重和层次结构重建的。此外，重建了一个“根级别”的 forward 函数来捕获这些分区之间的数据流。这些数据流稍后将以分布式的方式由流水线运行时重放。
+
+`Pipe` 对象提供了一个用于检索“模型分区”的方法：
+
+返回的 `stage_mod` 是一个 `nn.Module`，你可以使用它来创建优化器、保存或加载检查点，或应用其他并行机制。
+
+`Pipe` 还允许你在给定 `ProcessGroup` 的情况下在设备上创建分布式阶段运行时：
+
+或者，如果你想在修改 `stage_mod` 之后再构建阶段运行时，你可以使用函数版本的 `build_stage` API。例如：
+
+流水线前端使用跟踪器（`torch.export`）将你的模型捕获到一个单一的图中。如果你的模型无法进行全图捕获，你可以使用下面的手动前端。
+
+在最初创建此包的 PiPPy 存储库中，我们保留了基于未修改的 Hugging Face 模型的示例。请参阅 `examples/huggingface` 目录。
+
+首先，流水线 API 通过跟踪模型将其转换为有向无环图（DAG）。它使用 `torch.export`——一种 PyTorch 2 全图捕获工具来跟踪模型。
+
+然后，它将一个阶段所需的操作和参数组合在一起，重建为一个子模块：`submod_0`，`submod_1`，...
+
+与传统的子模块访问方法（如 `Module.children()`）不同，流水线 API 不仅会切断模型的结构，还会切断模型的 forward 函数。
+
+这是必要的，因为像 `Module.children()` 这样的模型结构仅在 `Module.__init__()` 期间捕获信息，而不捕获有关 `Module.forward()` 的任何信息。换句话说，`Module.children()` 缺少有关以下对流水线至关重要的方面的信息：
+
+子模块在 forward 中的执行顺序
+
+子模块之间的激活流
+
+子模块之间是否存在函数式运算符（例如，`relu` 或 `add` 操作不会被 `Module.children()` 捕获）。
+
+相反，流水线 API 确保真正保留了 forward 行为。它还捕获分区之间的激活流，帮助分布式运行时在无需人工干预的情况下发出正确的发送/接收调用。
+
+流水线 API 的另一个灵活性在于，分割点可以位于模型层次结构中的任意级别。在分割的分区中，与该分区相关的原始模型层次结构将为你重建，不产生额外成本。结果，指向子模块或参数的完全限定名（FQN）仍然有效，依赖于 FQN 的服务（如 FSDP、TP 或检查点）可以在你的分区模块上运行，几乎不需要更改代码。
+
+你可以通过扩展以下两个类之一来实现自己的流水线调度：
+
+`PipelineScheduleSingle`
+
+`PipelineScheduleMulti`
+
+`PipelineScheduleSingle` 用于每个 rank 仅分配一个阶段的调度。`PipelineScheduleMulti` 用于每个 rank 分配多个阶段的调度。
+
+例如，`ScheduleGPipe` 和 `Schedule1F1B` 是 `PipelineScheduleSingle` 的子类。而 `ScheduleInterleaved1F1B`、`ScheduleLoopedBFS`、`ScheduleInterleavedZeroBubble` 和 `ScheduleZBVZeroBubble` 是 `PipelineScheduleMulti` 的子类。
+
+你可以使用 `torch._logging` 中的 `TORCH_LOGS` 环境变量开启附加日志记录：
+
+`TORCH_LOGS=+pp` 将显示 `logging.DEBUG` 消息及高于它的所有级别的消息。
+
+`TORCH_LOGS=pp` 将显示 `logging.INFO` 消息及更高级别的消息。
+
+`TORCH_LOGS=-pp` 将显示 `logging.WARNING` 消息及更高级别的消息。
+
+下面的这组 API 将你的模型转换为流水线表示。
+
+表示子模块执行过程中可以发生分割的点的枚举。 :ivar BEGINNING: 表示在 forward 函数中某个子模块执行之前添加分割点。 :ivar END: 表示在 forward 函数中某个子模块执行之后添加分割点。
+
+基于规范拆分模块。
+
+有关更多详细信息，请参见 `Pipe`。
+
+module (Module) – 要拆分的模块。
+
+mb_args (tuple[Any, ...]) – 微批次形式的示例位置参数输入。
+
+mb_kwargs (可选[dict[str, Any]]) – 微批次形式的示例关键字参数输入。（默认值：None）
+
+split_spec (可选[dict[str, torch.distributed.pipelining._IR.SplitPoint]]) – 使用子模块名称作为分割标记的字典。（默认值：None）
+
+split_policy (可选[Callable[[GraphModule], GraphModule]]) – 用于拆分模块的策略。（默认值：None）
+
+`Pipe` 类的流水线表示。
+
+`pipe_split` 是一个特殊的运算符，用于标记模块中阶段之间的边界。它用于将模块划分为多个阶段。如果你标记的模块在 eager 模式下运行，它是一个空操作。
+
+上面的示例将被分割成两个阶段。
+
+用于指定输入分块的类
+
+根据它们各自的分块规范，将给定的 args 和 kwargs 序列分割成若干块。
+
+args (tuple[Any, ...]) – args 组成的元组
+
+kwargs (可选[dict[str, Any]]) – kwargs 组成的字典
+
+chunks (int) – 将 args 和 kwargs 切分成的块数
+
+args_chunk_spec (可选[tuple[torch.distributed.pipelining.microbatch.TensorChunkSpec, ...]]) – args 的分块规范，形状与 args 相同
+
+kwargs_chunk_spec (可选[dict[str, torch.distributed.pipelining.microbatch.TensorChunkSpec]]) – kwargs 的分块规范，形状与 kwargs 相同
+
+分片后的 args 列表 kwargs_split: 分片后的 kwargs 列表
+
+根据分块规范将给定的块列表合并为单个值。
+
+chunks (list[Any]) – 块组成的列表
+
+chunk_spec – 各个块的分布规范
+
+表示流水线并行设置中一个流水线阶段的类。
+
+`PipelineStage` 假设模型进行顺序分区，即模型被分成多个块，其中一个块的输出馈送到下一个块的输入中，没有跳跃连接。
+
+`PipelineStage` 通过将输出从 stage0 线性顺序地传播到 stage1 等，自动执行运行时形状/dtype 推断。要绕过形状推断，请将 `input_args` 和 `output_args` 传递给每个 `PipelineStage` 实例。
+
+submodule (nn.Module) – 被此阶段包装的 PyTorch 模块。
+
+stage_index (int) – 此阶段的 ID。
+
+num_stages (int) – 阶段总数。
+
+device (torch.device) – 此阶段所在的设备。
+
+input_args (Union[torch.Tensor, Tuple[torch.tensor]], 可选) – 子模块的输入参数。
+
+output_args (Union[torch.Tensor, Tuple[torch.tensor]], 可选) – 子模块的输出参数。
+
+group (dist.ProcessGroup, 可选) – 用于分布式训练的进程组。如果为 None，则使用默认组。
+
+dw_builder (可选[Callable[[], Callable[..., None]]) – 如果提供，`dw_builder` 将构建一个新的 `dw_runner` 函数，用于在 F、I、W（前向、输入、权重）零气泡调度中执行 W 操作（输入权重）。
+
+在给定要被此阶段包装的 `stage_module` 和流水线信息的情况下创建一个流水线阶段。
+
+stage_module (torch.nn.Module) – 被此阶段包装的模块
+
+stage_index (int) – 此阶段在流水线中的索引
+
+pipe_info (PipeInfo) – 有关流水线的信息，可以通过 `pipe.info()` 检索
+
+device (torch.device) – 此阶段使用的设备
+
+group (可选[dist.ProcessGroup]) – 此阶段使用的进程组
+
+一个可与 `PipelineSchedules` 一起运行的流水线阶段。
+
+GPipe 调度。将以填充-排空的方式处理所有微批次。
+
+1F1B 调度。将在稳定状态下对微批次执行一次前向传播和一次反向传播。
+
+Interleaved 1F1B 调度。有关详细信息，请参见 https://arxiv.org/pdf/2104.04473。将在稳定状态下对微批次执行一次前向传播和一次反向传播，并支持每个 rank 对应多个阶段。当多个本地阶段的微批次准备就绪时，Interleaved 1F1B 会优先处理较早的微批次（也称为“深度优先”）。
+
+此调度与原论文大致相同。区别在于它放宽了对 `num_microbatch % pp_size == 0` 的要求。使用 `flex_pp` 调度，我们将得到 `num_rounds = max(1, n_microbatches // pp_group_size)`，只要 `n_microbatches % num_rounds` 为 0，它就能正常工作。举几个支持的例子：
+
+pp_group_size = 4，n_microbatches = 10。我们会有 num_rounds = 2 并且 n_microbatches % 2 为 0。
+
+pp_group_size = 4，n_microbatches = 3。我们会有 num_rounds = 1 并且 n_microbatches % 1 为 0。
+
+广度优先流水线并行。有关详细信息，请参见 https://arxiv.org/abs/2211.05953。与 Interleaved 1F1B 类似，Looped BFS 支持每个 rank 对应多个阶段。不同之处在于，当多个本地阶段的微批次准备就绪时，Loops BFS 将优先处理较早的阶段，立即运行所有可用的微批次。
+
+Interleaved Zero Bubble 调度。有关详细信息，请参见 https://arxiv.org/pdf/2401.10241。将在稳定状态下对微批次的输入执行一次前向传播和一次反向传播，并支持每个 rank 对应多个阶段。使用权重的反向传播来填补流水线气泡。
+
+具体而言，这是在实现论文中的 ZB1P 调度。
+
+Zero Bubble 调度（ZBV 变体）。有关详细信息，请参见 https://arxiv.org/pdf/2401.10241 第 6 节。
+
+此调度要求每个 rank 恰好有两个阶段。
+
+此调度将在稳定状态下对微批次的输入执行一次前向传播和一次反向传播，并支持每个 rank 对应多个阶段。使用关于权重的反向传播来填补流水线气泡。
+
+只有当前向时间 == 输入反向时间 == 权重反向时间时，此 ZB-V 调度才具有“零气泡”属性。在实践中，这对于实际模型不太可能是真实的，因此或者可以为不相等/不平衡的时间实现一个贪心调度器。
+
+DualPipeV 调度。基于 DeepSeek 在 https://arxiv.org/pdf/2412.19437 中引入的 DualPipe 调度的更高效的调度变体
+
+基于 deepseek-ai/DualPipe 的开源代码
+
+单阶段调度的基类。实现了 `step` 方法。派生类应该实现 `_step_microbatches`。
+
+根据 `scale_grads` 参数（默认为 True），梯度会按 `num_microbatches` 进行缩放。此设置应与你的 `loss_fn` 配置相匹配，它可以是对损失求平均值（`scale_grads=True`）或对损失求和（`scale_grads=False`）。
+
+使用整批输入运行一次流水线调度的迭代。会自动将输入切分成微批次，并根据调度实现来处理这些微批次。
+
+args：模型的位置参数（与非流水线情况一样）。kwargs：模型的关键字参数（与非流水线情况一样）。target：损失函数的目标值。losses：一个用于存储每个微批次损失的列表。
+
+多阶段调度的基类。实现了 `step` 方法。
+
+根据 `scale_grads` 参数（默认为 True），梯度会按 `num_microbatches` 进行缩放。此设置应与你的 `loss_fn` 配置相匹配，它可以是对损失求平均值（`scale_grads=True`）或对损失求和（`scale_grads=False`）。
+
+使用整批输入运行一次流水线调度的迭代。会自动将输入切分成微批次，并根据调度实现来处理这些微批次。
+
+args：模型的位置参数（与非流水线情况一样）。kwargs：模型的关键字参数（与非流水线情况一样）。target：损失函数的目标值。losses：一个用于存储每个微批次损失的列表。
+
+---
+
+## 张量并行 - torch.distributed.tensor.parallel#
+
+**URL:** https://pytorch.org/docs/stable/distributed.tensor.parallel.html
+
+**目录：**
+- 张量并行 - torch.distributed.tensor.parallel#
+
+创建于：2025年6月13日 | 最后更新于：2025年6月13日
+
+张量并行（TP）构建在 PyTorch DistributedTensor (DTensor)[https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/README.md] 之上，提供了不同的并行方式：Colwise、Rowwise 和 Sequence Parallelism。
+
+张量并行 API 是实验性的，随时可能更改。
+
+使用张量并行来并行化你的 `nn.Module` 的入口点是：
+
+通过基于用户指定的计划并行化模块或子模块，在 PyTorch 中应用张量并行。
+
+我们基于 `parallelize_plan` 并行化模块或子模块。`parallelize_plan` 包含 `ParallelStyle`，指示用户希望如何对模块或子模块进行并行化。
+
+用户还可以通过模块的完全限定名（FQN）为每个模块指定不同的并行方式。
+
+请注意，`parallelize_module` 仅接受一维的 `DeviceMesh`。如果你有二维或 N 维的 `DeviceMesh`，请首先将 `DeviceMesh` 切分为一维的子 `DeviceMesh`，然后再传递给此 API（即 `device_mesh["tp"]`）。
+
+module (nn.Module) – 要并行化的模块。
+
+device_mesh (DeviceMesh, 可选) – 描述 DTensor 的设备网格拓扑的对象。如果未指定，则该调用必须在 `DeviceMesh` 上下文中。
+
+parallelize_plan (Union[ParallelStyle, Dict[str, ParallelStyle]], 可选) – 用于并行化模块的计划。它可以是包含我们如何为张量并行准备输入/输出的 `ParallelStyle` 对象，也可以是模块 FQN 及其相应 `ParallelStyle` 对象的字典。如果未指定，该调用此刻将不执行任何操作。
+
+src_data_rank (int, 可选) – 逻辑/全局张量的源数据的 rank，它被 `distribute_tensor()` 用于向其他 rank 分发/广播分片/副本。默认情况下，我们在每个 `DeviceMesh` 维度上使用 `group_rank=0` 作为源数据，以保留单设备语义。如果显式传入 None，`parallelize_module()` 将仅使用其本地数据，而不是尝试通过分发/广播来保留单设备语义。默认：0
+
+一个被并行化的 `nn.Module` 对象。
+
+对于复杂的模块架构（如 Attention、MLP 层），我们建议将不同的 `ParallelStyle`（即 `ColwiseParallel` 和 `RowwiseParallel`）组合在一起，并作为 `parallelize_plan` 传入，以实现所需的分片计算。
+
+张量并行支持以下并行方式：
+
+按列方式对兼容的 `nn.Module` 进行分区。目前支持 `nn.Linear` 和 `nn.Embedding`。用户可以将其与 `RowwiseParallel` 组合在一起，实现更复杂模块（即 MLP、Attention）的分片。
+
+input_layouts (Placement, 可选) – `nn.Module` 输入张量的 DTensor 布局，用于标注输入张量使其成为 DTensor。如果未指定，我们假定输入张量是复制。
+
+output_layouts (Placement, 可选) – `nn.Module` 输出的 DTensor 布局，用于确保 `nn.Module` 的输出具有用户期望的布局。如果未指定，输出张量会在最后一个维度上进行分片。
+
+use_local_output (bool, 可选) – 是否为模块输出使用本地 `torch.Tensor` 而不是 DTensor，默认值：True。
+
+一个表示 `nn.Module` 的 Colwise 分片的 `ParallelStyle` 对象。
+
+默认情况下，如果未指定 `output_layouts`，`ColwiseParallel` 的输出会在最后一个维度上进行分片，如果有运算符需要特定的张量形状（即在配对的 `RowwiseParallel` 之前），请记住，如果输出被分片了，可能需要对运算符进行调整以适应分片后的大小。
+
+按行方式对兼容的 `nn.Module` 进行分区。目前支持 `nn.Linear` 和 `nn.Embedding`。用户可以将其与 `ColwiseParallel` 组合在一起，实现更复杂模块（即 MLP、Attention）的分片。
+
+input_layouts (Placement, 可选) – `nn.Module` 输入张量的 DTensor 布局，用于标注输入张量使其成为 DTensor。如果未指定，我们假定输入张量在最后一个维度上进行分片。
+
+output_layouts (Placement, 可选) – `nn.Module` 输出的 DTensor 布局，用于确保 `nn.Module` 的输出具有用户期望的布局。如果未指定，输出张量为复制。
+
+use_local_output (bool, 可选) – 是否为模块输出使用本地 `torch.Tensor` 而不是 DTensor，默认值：True。
+
+一个表示 `nn.Module` 的 Rowwise 分片的 `ParallelStyle` 对象。
+
+SequenceParallel 会复制兼容的 `nn.Module` 参数，并运行输入在序列维度上进行分片后的计算。目前支持 `nn.LayerNorm`、`nn.Dropout` 和 RMSNorm 的 python 实现。
+
+此方式实现了论文《Reducing Activation Recomputation in Large Transformer Models》中描述的操作。
+
+如果传入此 `nn.Module` 的输入是 `torch.Tensor`，它会假定输入已经在序列维度上进行了分片，并将输入转换为在序列维度上分片的 DTensor。如果传入此 `nn.Module` 的输入已经是 DTensor 但并未在序列维度上分片，它将重新分布输入，以实现在序列维度上分片。
+
+`nn.Module` 的输出将在序列维度上分片。
+
+sequence_dim (int, 可选) – `nn.Module` 的输入张量的序列维度，用于标注输入张量使其成为在序列维度上分片的 DTensor，默认值：1。
+
+use_local_output (bool, 可选) – 是否为模块输出使用本地 `torch.Tensor` 而不是 DTensor，默认值：False。
+
+一个表示 `nn.Module` 的 Sequence Parallel 的 `ParallelStyle` 对象。
+
+`SequenceParallel` 方式假定 `nn.Module` 中有权重时（如 `nn.LayerNorm` 或 RMSNorm）进行全 1 初始化（且它们默认就是全 1 初始化）。如果这些模块上的权重有自定义初始化，则需要在并行化前后广播权重，以确保它们被复制。
+
+为了仅使用 DTensor 布局配置 `nn.Module` 的输入和输出并执行必要的布局重分布，而不将模块参数分配为 DTensor，可以在调用 `parallelize_module` 时在 `parallelize_plan` 中使用以下 `ParallelStyle`：
+
+根据 `input_layouts` 将 `nn.Module` 的输入张量在运行时转换为 DTensor，并根据 `desired_input_layouts` 执行布局重分布，从而配置 `nn.Module` 的输入。
+
+input_layouts (Union[Placement, Tuple[可选[Placement]]]) – `nn.Module` 的输入张量的 DTensor 布局，用于将输入张量转换为 DTensor。如果某些输入不是 `torch.Tensor` 或无需转换为 DTensor，则需要指定 None 作为占位符。默认值：None。
+
+desired_input_layouts (Union[Placement, Tuple[可选[Placement]]]) – `nn.Module` 输入张量的所需 DTensor 布局，用于确保 `nn.Module` 的输入具有所需的 DTensor 布局。此参数需要与 `input_layouts` 具有相同的长度。默认值：None。
+
+input_kwarg_layouts (Dict[str, Placement]) – `nn.Module` 的输入 kwargs 的 DTensor 布局，用于将输入 kwarg 张量转换为 DTensor。默认值：None
+
+desired_input_kwarg_layouts – (Dict[str, Placement])：`nn.Module` 的输入 kwargs 的所需 DTensor 布局，用于确保 `nn.Module` 的输入具有所需的 DTensor 布局。默认值：None。
+
+use_local_output (bool, 可选) – 是否为模块输入使用本地 `torch.Tensor` 而不是 DTensor，默认值：False。
+
+一个用于准备 `nn.Module` 输入分片布局的 `ParallelStyle` 对象。
+
+根据 `output_layouts` 将 `nn.Module` 的输出张量在运行时转换为 DTensor，并根据 `desired_output_layouts` 执行布局重分布，从而配置 `nn.Module` 的输出。
+
+output_layouts (Union[Placement, Tuple[Placement]]) – `nn.Module` 的输出张量的 DTensor 布局，如果它们是 `torch.Tensor`，则用于将其转换为 DTensor。如果某些输出不是 `torch.Tensor` 或无需转换为 DTensor，则需要指定 None 作为占位符。
+
+desired_output_layouts (Union[Placement, Tuple[Placement]]) – `nn.Module` 的输出张量的所需 DTensor 布局，用于确保 `nn.Module` 的输出具有所需的 DTensor 布局。
+
+use_local_output (bool, 可选) – 是否为模块输出使用本地 `torch.Tensor` 而不是 DTensor，默认值：True。
+
+一个用于准备 `nn.Module` 输出分片布局的 `ParallelStyle` 对象。
+
+根据 `input_layouts`（和 `output_layouts`）将 `nn.Module` 的输入张量（以及相应的输出张量）在运行时转换为 DTensor，并根据 `desired_input_layouts`（以及 `desired_output_layouts`）执行布局重分布，从而配置 `nn.Module` 的输入（和输出）。这是 `PrepareModuleInput` 和 `PrepareModuleOutput` 的组合。
+
+input_layouts (Union[Placement, Tuple[可选[Placement]]]) – `nn.Module` 的输入张量的 DTensor 布局，用于将输入张量转换为 DTensor。如果某些输入不是 `torch.Tensor` 或无需转换为 DTensor，则需要指定 None 作为占位符。默认值：None。
+
+desired_input_layouts (Union[Placement, Tuple[可选[Placement]]]) – `nn.Module` 输入张量的所需 DTensor 布局，用于确保 `nn.Module` 的输入具有所需的 DTensor 布局。此参数需要与 `input_layouts` 具有相同的长度。默认值：None。
+
+input_kwarg_layouts (Dict[str, Placement]) – `nn.Module` 的输入 kwargs 的 DTensor 布局，用于将输入 kwarg 张量转换为 DTensor。默认值：None
+
+desired_input_kwarg_layouts – (Dict[str, Placement])：`nn.Module` 的输入 kwargs 的所需 DTensor 布局，用于确保 `nn.Module` 的输入具有所需的 DTensor 布局。默认值：None。
+
+use_local_input (bool, 可选) – 是否为模块输入使用本地 `torch.Tensor` 而不是 DTensor，默认值：False。
+
+output_layouts (Union[Placement, Tuple[Placement]]) – `nn.Module` 的输出张量的 DTensor 布局，如果它们是 `torch.Tensor`，则用于将其转换为 DTensor。如果某些输出不是 `torch.Tensor` 或无需转换为 DTensor，则需要指定 None 作为占位符。
+
+desired_output_layouts (Union[Placement, Tuple[Placement]]) – `nn.Module` 的输出张量的所需 DTensor 布局，用于确保 `nn.Module` 的输出具有