原文：
译者：飞龙
协议：CC BY-NC-SA 4.0
作者：Hamid Shojanazeri，Yanli Zhao，Shen Li
注意
在github上查看并编辑本教程。
在大规模训练 AI 模型是一项具有挑战性的任务，需要大量的计算能力和资源。同时，处理这些非常大模型的训练也伴随着相当大的工程复杂性。PyTorch FSDP，在 PyTorch 1.11 中发布，使这变得更容易。
在本教程中，我们展示了如何使用FSDP APIs，用于简单的 MNIST 模型，可以扩展到其他更大的模型，比如HuggingFace BERT 模型，GPT 3 模型高达 1T 参数。示例 DDP MNIST 代码是从这里借鉴的。
在DistributedDataParallel（DDP）训练中，每个进程/工作器拥有模型的副本并处理一批数据，最后使用全局归约来汇总不同工作器上的梯度。在 DDP 中，模型权重和优化器状态在所有工作器之间复制。FSDP 是一种数据并行 ism，它在 DDP 等级之间分片模型参数、优化器状态和梯度。
使用 FSDP 进行训练时，GPU 内存占用比在所有工作节点上使用 DDP 进行训练时要小。这使得一些非常大的模型的训练变得可行，因为可以容纳更大的模型或批量大小在设备上。但这也会增加通信量。通过内部优化，如重叠通信和计算，可以减少通信开销。
FSDP 工作流程
在高层次上，FSDP 的工作方式如下：
在构造函数中
- 分片模型参数和每个等级只保留自己的分片
在前向路径
- 运行 all_gather 来收集所有等级的所有碎片，以恢复此 FSDP 单元中的完整参数。
- 进行前向计算
- 丢弃刚收集的参数分片
在反向路径中
- 运行 all_gather 来收集所有等级的所有碎片，以恢复此 FSDP 单元中的完整参数。
- 运行反向计算
- 运行 reduce_scatter 来同步梯度
- 丢弃参数。
将 FSDP 的分片视为将 DDP 梯度全局归约分解为归约散射和全局聚集的一种方式。具体来说，在反向传播过程中，FSDP 减少并散射梯度，确保每个秩具有梯度的一个片段。然后在优化器步骤中更新相应的参数片段。最后，在随后的前向传播过程中，它执行全局聚集操作来收集和组合更新的参数片段。
FSDP Allreduce
在这里，我们使用一个玩具模型来对 MNIST 数据集进行训练，以演示目的。这些 API 和逻辑也可以应用于训练更大的模型。
设置
1.1 安装 PyTorch 和 Torchvision
我们将以下代码片段添加到一个名为“FSDP_mnist.py”的 Python 脚本中。
1.2 导入必要的包
注意
本教程适用于 PyTorch 版本 1.12 及更高版本。如果您使用的是早期版本，请将所有的 size_based_auto_wrap_policy 实例替换为 default_auto_wrap_policy。
1.3 分布式训练设置。正如我们提到的，FSDP 是一种数据并行 ism，它需要一个分布式训练环境，因此我们在这里使用两个辅助函数来初始化分布式训练的进程并进行清理。
2.1 定义我们的手写数字分类的玩具模型。
2.2 定义一个训练函数。
2.3 定义一个验证函数
2.4 定义一个分布式训练函数，将模型包装在 FSDP 中
注意：为了保存 FSDP 模型，我们需要在每个排名上调用 state_dict，然后在排名 0 上保存整体状态。
2.5 最后，解析参数并设置主函数
我们已记录了 CUDA 事件来测量 FSDP 模型特定部分的时间。CUDA 事件时间为 110.85 秒。
使用 FSDP 包装模型后，模型将如下所示，我们可以看到模型已经被包装在一个 FSDP 单元中。或者，我们将考虑接下来添加 fsdp_auto_wrap_policy，并讨论其中的区别。
以下是在 g4dn.12.xlarge AWS EC2 实例上使用 4 个 GPU 进行 FSDP MNIST 训练时从 PyTorch Profiler 捕获的峰值内存使用情况。
FSDP 峰值内存使用量
在 FSDP 中应用fsdp_auto_wrap_policy，否则，FSDP 将把整个模型放在一个 FSDP 单元中，这将降低计算效率和内存效率。它的工作方式是，假设您的模型包含 100 个线性层。如果您对模型进行 FSDP 处理，那么只会有一个包含整个模型的 FSDP 单元。在这种情况下，allgather 将收集所有 100 个线性层的完整参数，因此不会为参数分片节省 CUDA 内存。此外，对于所有 100 个线性层，只有一个阻塞的 allgather 调用，层之间不会有通信和计算重叠。
为了避免这种情况，您可以传入一个 fsdp_auto_wrap_policy，当满足指定条件（例如大小限制）时，将封装当前的 FSDP 单元并自动启动一个新的。这样，您将拥有多个 FSDP 单元，每次只需要一个 FSDP 单元收集完整参数。例如，假设您有 5 个 FSDP 单元，每个包含 20 个线性层。然后，在前向传播中，第一个 FSDP 单元将收集前 20 个线性层的参数，进行计算，丢弃参数，然后继续下一个 20 个线性层。因此，在任何时候，每个 rank 只会实现 20 个线性层的参数/梯度，而不是 100 个。
为了在 2.4 中这样做，我们定义了 auto_wrap_policy 并将其传递给 FSDP 包装器，在以下示例中，my_auto_wrap_policy 定义了如果该层中的参数数量大于 100，则该层可以被 FSDP 包装或分片。如果该层中的参数数量小于 100，则它将与其他小层一起被 FSDP 包装。找到一个最佳的自动包装策略是具有挑战性的，PyTorch 将在将来为此配置添加自动调整功能。没有自动调整工具，最好通过实验使用不同的自动包装策略来分析您的工作流程，并找到最佳策略。
应用 fsdp_auto_wrap_policy，模型将如下：
以下是在 g4dn.12.xlarge AWS EC2 实例上使用 4 个 GPU 进行 MNIST 训练时，从 PyTorch Profiler 捕获的 FSDP 自动包装策略的峰值内存使用情况。可以观察到，与未应用自动包装策略的 FSDP 相比，每个设备上的峰值内存使用量较小，从约 75 MB 降至 66 MB。
使用 Auto_wrap 策略的 FSDP 峰值内存使用量
CPU 卸载：如果模型非常庞大，即使使用 FSDP 也无法适应 GPU，那么 CPU 卸载可能会有所帮助。
目前，仅支持参数和梯度的 CPU 卸载。可以通过传入 cpu_offload=CPUOffload(offload_params=True)来启用。
请注意，目前这将隐式地启用梯度卸载到 CPU，以便参数和梯度在同一设备上与优化器一起工作。此 API 可能会发生变化。默认值为 None，在这种情况下将不会进行卸载。
使用这个功能可能会显著减慢训练速度，因为频繁地从主机复制张量到设备，但它可以帮助提高内存效率并训练更大规模的模型。
在 2.4 版本中，我们只是将其添加到 FSDP 包装器中。
将其与 DDP 进行比较，如果在 2.4 中我们只是正常地将模型包装在 DPP 中，并保存更改在“DDP_mnist.py”中。
以下是在 g4dn.12.xlarge AWS EC2 实例上使用 4 个 GPU 进行 DDP MNIST 训练时从 PyTorch 分析器中捕获的峰值内存使用情况。
DDP 使用 Auto_wrap 策略的峰值内存使用量
考虑到我们在这里定义的玩具示例和微小的 MNIST 模型，我们可以观察到 DDP 和 FSDP 的峰值内存使用之间的差异。在 DDP 中，每个进程持有模型的副本，因此内存占用量较高，而与 DDP 排名相比，FSDP 将模型参数、优化器状态和梯度进行分片。使用 auto_wrap 策略的 FSDP 的峰值内存使用量最低，其次是 FSDP 和 DDP。
此外，从时间上看，考虑到小模型并在单台机器上运行训练，FSDP 在有或没有自动包装策略的情况下几乎与 DDP 一样快。这个例子并不代表大多数真实应用程序，有关 DDP 和 FSDP 的详细分析和比较，请参考这篇博客文章。
原文：
译者：飞龙
协议：CC BY-NC-SA 4.0
作者：Hamid Shojanazeri，Less Wright，Rohan Varma，Yanli Zhao
本教程介绍了 PyTorch 1.12 版本中 Fully Sharded Data Parallel（FSDP）的更高级特性。要熟悉 FSDP，请参考FSDP 入门教程。
在本教程中，我们使用 FSDP 对 HuggingFace（HF）的 T5 模型进行微调，作为文本摘要的工作示例。
这个例子使用了 Wikihow，为了简单起见，我们将展示在一个单节点上进行训练，使用带有 8 个 A100 GPU 的 P4dn 实例。我们很快将在多节点集群上发布一篇关于大规模 FSDP 训练的博客文章，请关注 PyTorch 的官方媒体渠道。
FSDP 是一个已经准备好的软件包，专注于易用性、性能和长期支持。FSDP 的主要优势之一是减少每个 GPU 上的内存占用。这使得可以使用更低的总内存训练更大的模型，同时利用计算和通信的重叠来高效训练模型。这种减少的内存压力可以用来训练更大的模型或增加批量大小，潜在地帮助提高整体训练吞吐量。您可以在这里阅读更多关于 PyTorch FSDP 的信息。
- Transformer 自动包装策略
- 混合精度
- 在设备上初始化 FSDP 模型
- 分片策略
- 向后预取
- 通过流式传输保存模型检查点到 CPU
在高层次上，FDSP 的工作方式如下：
在构造函数中
- 分片模型参数，每个等级只保留自己的分片
在前向传播中
- 运行 all_gather 以收集所有排名的所有碎片，以恢复此 FSDP 单元的完整参数 运行前向计算
- 丢弃刚刚收集的非所有者参数分片以释放内存
在反向传递中
- 运行 all_gather 以收集所有等级的所有碎片，以恢复此 FSDP 单元中的完整参数 运行向后计算
- 丢弃非所有者参数以释放内存。
- 运行 reduce_scatter 以同步梯度。
HF T5 预训练模型有四种不同大小可供选择，从参数为 6000 万的小型模型到参数为 110 亿的 XXL 模型。在本教程中，我们演示了使用 WikiHow 数据集对 T5 3B 进行微调，以用于文本摘要。本教程的主要重点是突出 FSDP 中可用的不同功能，这些功能有助于训练超过 3B 参数的大规模模型。此外，我们还介绍了基于 Transformer 的模型的特定功能。本教程的代码可在Pytorch 示例中找到。
设置
1.1 安装 PyTorch 最新版本
我们将安装 PyTorch 的 nightlies 版本，因为一些功能，比如激活检查点，在 nightlies 版本中可用，并将在 1.12 版本之后的下一个 PyTorch 发布中添加。
1.2 数据集设置
请创建一个名为 data 的文件夹，从wikihowAll.csv和wikihowSep.cs下载 WikiHow 数据集，并将它们放在 data 文件夹中。我们将使用来自summarization_dataset的 wikihow 数据集。
接下来，我们将以下代码片段添加到一个名为“T5_training.py”的 Python 脚本中。
注意
本教程的完整源代码可在PyTorch 示例中找到。
1.3 导入必要的包：
1.4 分布式训练设置。在这里，我们使用两个辅助函数来初始化分布式训练的进程，然后在训练完成后进行清理。在本教程中，我们将使用 torch elastic，使用torchrun，它会自动设置工作进程的 RANK 和 WORLD_SIZE。
2.1 设置 HuggingFace T5 模型：
我们还在这里添加了一些用于日期和格式化内存指标的辅助函数。
2.2 定义一个训练函数：
2.3 定义一个验证函数：
定义一个包装模型在 FSDP 中的分布式训练函数。
2.5 解析参数并设置主函数：
使用 torchrun 运行训练：
正如在上一个教程中讨论的，auto_wrap_policy 是 FSDP 功能之一，它使得自动对给定模型进行分片并将模型、优化器和梯度分片放入不同的 FSDP 单元变得容易。
对于一些架构，比如 Transformer 编码器-解码器，模型的一些部分，比如嵌入表，被编码器和解码器共享。在这种情况下，我们需要将嵌入表放在外部 FSDP 单元中，以便从编码器和解码器中访问。此外，通过为 transformer 注册层类，分片计划可以变得更加通信高效。在 PyTorch 1.12 中，FSDP 添加了这种支持，现在我们有了一个用于 transformers 的包装策略。
可以按照以下方式创建，其中 T5Block 代表 T5 变压器层类（包含 MHSA 和 FFN）。
要查看包装的模型，您可以轻松打印模型并直观地检查分片和 FSDP 单元。
FSDP 支持灵活的混合精度训练，允许使用任意降低精度类型（如 fp16 或 bfloat16）。目前，BFloat16 仅在安培 GPU 上可用，因此在使用之前需要确认是否有本机支持。例如，在 V100 上，仍然可以运行 BFloat16，但由于它是非本机运行，可能会导致显著的减速。
要检查是否原生支持 BFloat16，您可以使用以下方法：
在 FSDP 中混合精度的一个优点是为参数、梯度和缓冲区提供不同精度级别的细粒度控制。
请注意，如果某种类型（参数、减少、缓冲区）未指定，则它们将不会被转换。
这种灵活性使用户可以进行精细的控制，比如只将梯度通信设置为以降低精度进行，而所有参数/缓冲计算则以全精度进行。在节点内通信是主要瓶颈且参数/缓冲必须以全精度进行以避免精度问题的情况下，这种方法可能非常有用。可以使用以下策略来实现：
在 2.4 版本中，我们只需将相关的混合精度策略添加到 FSDP 包装器中：
在我们的实验中，我们观察到使用 BFloat16 进行训练可以加快速度达到 4 倍，并且在一些实验中可以减少大约 30%的内存，这可以用于增加批量大小。
在 1.12 版本中，FSDP 支持一个 device_id 参数，旨在初始化设备上的输入 CPU 模块。当整个模型无法适应单个 GPU，但适应主机的 CPU 内存时，这将非常有用。当指定 device_id 时，FSDP 将根据每个 FSDP 单元将模型移动到指定的设备上，避免 GPU 内存不足问题，同时初始化速度比基于 CPU 的初始化快数倍。
默认情况下，FSDP 分片策略被设置为完全分片模型参数，梯度和优化器状态在所有等级之间分片（也称为 Zero3 分片）。如果您希望使用 Zero2 分片策略，仅对优化器状态和梯度进行分片，FSDP 支持通过将分片策略传递给 FSDP 初始化来实现此功能，如下所示：“ShardingStrategy.SHARD_GRAD_OP”，而不是“ShardingStrategy.FULL_SHARD”。
这将减少 FSDP 中的通信开销，在这种情况下，在前向传播和反向传播后保持完整的参数。
在反向传播过程中，这样做可以节省一次全局聚合操作，从而减少通信量，但会增加内存占用。请注意，完整的模型参数会在反向传播结束时被释放，全局聚合操作将在下一次前向传播中进行。
后向预取设置控制了何时应请求下一个 FSDP 单元的参数。通过将其设置为 BACKWARD_PRE，下一个 FSDP 单元的参数可以在当前单元的计算开始之前开始请求并到达。这会重叠所有收集通信和梯度计算，可以增加训练速度，但会略微增加内存消耗。可以在 2.4 版本中的 FSDP 包装器中利用它。
backward_prefetch 有两种模式，BACKWARD_PRE 和 BACKWARD_POST。BACKWARD_POST 意味着直到当前 FSDP 单元处理完成之前，不会请求下一个 FSDP 单元的参数，从而最大限度地减少内存开销。在某些情况下，使用 BACKWARD_PRE 可以将模型训练速度提高 2-10%，对于更大的模型，速度提高更为显著。
使用 FULL_STATE_DICT 保存模型检查点，该保存方式与本地模型相同，PyTorch 1.12 提供了一些实用工具来支持保存更大的模型。
首先，可以指定一个 FullStateDictConfig，允许仅在 rank 0 上填充 state_dict 并转移到 CPU。
在使用这种配置时，FSDP 将会收集模型参数，逐个将其转移到 CPU 上，仅在 rank 0 上进行。当 state_dict 最终保存时，它只会在 rank 0 上填充，并包含 CPU 张量。这避免了对于大于单个 GPU 内存的模型可能出现的 OOM，并允许用户对模型进行检查点，其大小大致等于用户机器上可用的 CPU RAM。
这个功能可以按照以下方式运行：
在本教程中，我们介绍了 Pytorch 1.12 中可用的许多 FSDP 的新功能，并以 HF T5 作为运行示例。特别是对于变压器模型，使用适当的包装策略，以及混合精度和向后预取应该可以加快您的训练速度。此外，诸如在设备上初始化模型和通过流式传输到 CPU 保存检查点等功能应该有助于避免处理大型模型时的 OOM 错误。
我们正在积极努力为下一个版本的 FSDP 添加新功能。如果您有反馈、功能请求、问题或在使用 FSDP 时遇到问题，请随时通过在PyTorch Github 存储库中打开问题与我们联系。
原文：
译者：飞龙
协议：CC BY-NC-SA 4.0
作者：Howard Huang https://github.com/H-Huang，Feng Tian，Shen Li，Min Si
注意
在 github 上查看并编辑本教程。
先决条件:
- PyTorch 分布式概述
- PyTorch 集体通信包
- PyTorch Cpp 扩展
- 使用 PyTorch 编写分布式应用程序
本教程演示了如何实现一个自定义的并将其插入PyTorch 分布式包，使用cpp 扩展。当您需要为硬件定制专门的软件堆栈，或者想要尝试新的集体通信算法时，这将非常有帮助。
PyTorch 集体通信支持多种广泛采用的分布式训练功能，包括DistributedDataParallel，ZeroRedundancyOptimizer，FullyShardedDataParallel。为了使相同的集体通信 API 能够与不同的通信后端一起工作，分布式包将集体通信操作抽象为Backend类。不同的后端可以作为的子类使用首选的第三方库来实现。PyTorch 分布式带有三个默认后端，，和。然而，除了这三个后端之外，还有其他通信库（例如UCC，OneCCL），不同类型的硬件（例如TPU，Trainum）和新兴的通信算法（例如Herring，Reduction Server）。因此，分布式包提供了扩展 API 来允许定制集体通信后端。
以下 4 个步骤展示了如何在 Python 应用程序代码中实现一个虚拟的后端并使用它。请注意，本教程侧重于演示扩展 API，而不是开发一个功能完善的通信后端。因此，后端只涵盖了 API 的一个子集（和），并且只是将张量的值设置为 0。
第一步是实现一个子类，覆盖目标集体通信 API，并运行自定义通信算法。扩展还需要实现一个子类，作为通信结果的 future，并允许在应用代码中异步执行。如果扩展使用第三方库，可以在子类中包含头文件并调用库 API。下面的两个代码片段展示了和的实现。请查看dummy collectives存储库以获取完整的实现。
后端构造函数是从 Python 端调用的，因此扩展还需要向 Python 公开构造函数 API。这可以通过添加以下方法来实现。在这个例子中，和被实例化方法忽略，因为在这个虚拟实现中没有使用它们。然而，真实世界的扩展应该考虑使用来执行会合并支持参数。
现在，扩展源代码文件已经准备好。我们可以使用cpp extensions来构建它。为此，创建一个文件，准备路径和命令。然后调用来安装扩展。
如果扩展依赖于第三方库，您还可以在 cpp 扩展 API 中指定和。请参考torch ucc项目作为一个真实的例子。
安装完成后，您可以在调用init_process_group时方便地使用后端，就像它是一个内置后端一样。
我们可以根据后端来指定调度，方法是改变的参数。我们可以通过将后端参数指定为，将 CPU 张量的集体分发到后端，将 CUDA 张量的集体分发到后端。
要将所有张量发送到后端，我们可以简单地将指定为后端参数。
原文：
译者：飞龙
协议：CC BY-NC-SA 4.0
作者：Shen Li
注意
在github中查看和编辑本教程。
先决条件：
- PyTorch 分布式概述
- RPC API 文档
本教程使用两个简单示例演示如何使用torch.distributed.rpc包构建分布式训练，该包最初作为 PyTorch v1.4 中的实验性功能引入。这两个示例的源代码可以在PyTorch 示例中找到。
之前的教程，使用分布式数据并行开始和使用 PyTorch 编写分布式应用程序，描述了分布式数据并行，支持一种特定的训练范式，其中模型在多个进程中复制，并且每个进程处理输入数据的一个部分。有时，您可能会遇到需要不同训练范式的情况。例如：
- 在强化学习中，从环境中获取训练数据可能相对昂贵，而模型本身可能非常小。在这种情况下，可能有用的是并行运行多个观察者并共享单个代理。在这种情况下，代理在本地处理训练，但应用程序仍需要库来在观察者和训练者之间发送和接收数据。
- 您的模型可能太大，无法适应单台机器上的 GPU，因此需要一个库来将模型分割到多台机器上。或者您可能正在实现一个参数服务器训练框架，其中模型参数和训练器位于不同的机器上。
torch.distributed.rpc包可以帮助处理上述情况。在情况 1 中，RPC和RRef允许从一个工作进程发送数据到另一个工作进程，同时轻松引用远程数据对象。在情况 2 中，分布式自动求导和分布式优化器使得执行反向传播和优化器步骤就像是本地训练一样。在接下来的两个部分中，我们将使用一个强化学习示例和一个语言模型示例演示torch.distributed.rpc的 API。请注意，本教程的目标不是构建最准确或高效的模型来解决给定问题，而是展示如何使用torch.distributed.rpc包构建分布式训练应用程序。
本部分描述了使用 RPC 构建玩具分布式强化学习模型的步骤，以解决来自OpenAI Gym的 CartPole-v1 问题。策略代码大部分是从现有的单线程示例中借用的，如下所示。我们将跳过“策略”设计的细节，重点放在 RPC 的用法上。
我们准备展示观察者。在这个例子中，每个观察者都创建自己的环境，并等待代理的命令来运行一个剧集。在每个剧集中，一个观察者最多循环次迭代，在每次迭代中，它使用 RPC 将其环境状态传递给代理，并获得一个动作。然后将该动作应用于其环境，并从环境中获得奖励和下一个状态。之后，观察者使用另一个 RPC 向代理报告奖励。再次请注意，这显然不是最有效的观察者实现。例如，一个简单的优化可以是将当前状态和上一个奖励打包在一个 RPC 中，以减少通信开销。然而，目标是演示 RPC API 而不是构建 CartPole 的最佳求解器。因此，在这个例子中，让我们保持逻辑简单，将这两个步骤明确表示。
代理的代码稍微复杂一些，我们将其分解为多个部分。在这个例子中，代理既充当训练者又充当主控，它向多个分布式观察者发送命令来运行剧集，并在本地记录所有动作和奖励，这些将在每个剧集后的训练阶段中使用。下面的代码显示了构造函数，其中大多数行都在初始化各种组件。最后的循环在其他工作进程上远程初始化观察者，并在本地保存这些观察者的。代理将在稍后使用这些观察者的来发送命令。应用程序不需要担心的生命周期。每个的所有者维护一个引用计数映射来跟踪其生命周期，并保证只要有任何的活动用户，远程数据对象就不会被删除。有关详细信息，请参阅设计文档。
接下来，代理向观察者公开两个 API，用于选择动作和报告奖励。这些函数仅在代理上本地运行，但将通过 RPC 由观察者触发。
让我们在代理上添加一个函数，告诉所有观察者执行一个剧集。在这个函数中，首先创建一个列表来收集异步 RPC 的 futures，然后循环遍历所有观察者的来进行异步 RPC。在这些 RPC 中，代理还将自身的传递给观察者，以便观察者也可以在代理上调用函数。如上所示，每个观察者将向代理发起 RPC，这些是嵌套的 RPC。每个剧集结束后，和将包含记录的动作概率和奖励。
最后，在一个剧集结束后，代理需要训练模型，这在下面的函数中实现。这个函数中没有 RPC，它主要是从单线程的示例中借用的。因此，我们跳过描述其内容。
有了、和类，我们准备启动多个进程执行分布式训练。在这个例子中，所有进程都运行相同的函数，并使用排名来区分它们的角色。排名 0 始终是代理，所有其他排名都是观察者。代理通过反复调用和来充当主控，直到运行奖励超过环境指定的奖励阈值。所有观察者都 passively 等待代理的命令。代码由rpc.init_rpc和rpc.shutdown包装，分别初始化和终止 RPC 实例。更多细节请参阅API 页面。
以下是在 world_size=2 时进行训练时的一些示例输出。
在这个例子中，我们展示了如何使用 RPC 作为通信工具在工作器之间传递数据，以及如何使用 RRef 引用远程对象。当然，您可以直接在 和 API 之上构建整个结构，或者使用其他通信/RPC 库。然而，通过使用 torch.distributed.rpc，您可以获得本地支持，并在幕后持续优化性能。
接下来，我们将展示如何结合 RPC 和 RRef 与分布式自动求导和分布式优化器来执行分布式模型并行训练。
在本节中，我们使用一个 RNN 模型来展示如何使用 RPC API 构建分布式模型并行训练。示例 RNN 模型非常小，可以轻松适应单个 GPU，但我们仍将其层分布到两个不同的工作器上以演示这个想法。开发人员可以应用类似的技术将更大的模型分布到多个设备和机器上。
RNN 模型设计借鉴了 PyTorch 示例 仓库中的单词语言模型，其中包含三个主要组件，一个嵌入表，一个层和一个解码器。下面的代码将嵌入表和解码器包装成子模块，以便它们的构造函数可以传递给 RPC API。在子模块中，我们故意将层放在 GPU 上以涵盖使用情况。在 v1.4 中，RPC 始终在目标工作器上创建 CPU 张量参数或返回值。如果函数接受 GPU 张量，则需要显式将其移动到适当的设备上。
通过上述子模块，我们现在可以使用 RPC 将它们组合在一起创建一个 RNN 模型。在下面的代码中，代表参数服务器，它承载嵌入表和解码器的参数。构造函数使用remote API 在参数服务器上创建一个对象和一个对象，并在本地创建子模块。在前向传播过程中，训练器使用的来找到远程子模块，并通过 RPC 将输入数据传递给并获取查找结果。然后，它通过本地的层运行嵌入，最后使用另一个 RPC 将输出发送到子模块。通常，为了实现分布式模型并行训练，开发人员可以将模型划分为子模块，调用 RPC 远程创建子模块实例，并在必要时使用来找到它们。正如您在下面的代码中所看到的，它看起来非常类似于单机模型并行训练。主要区别是用 RPC 函数替换。
在介绍分布式优化器之前，让我们添加一个辅助函数来生成模型参数的 RRef 列表，这将被分布式优化器使用。在本地训练中，应用程序可以调用来获取所有参数张量的引用，并将其传递给本地优化器进行后续更新。然而，在分布式训练场景中，相同的 API 不起作用，因为一些参数存在于远程机器上。因此，分布式优化器不是接受参数列表，而是接受列表，每个模型参数都有一个，用于本地和远程模型参数。辅助函数非常简单，只需调用并在每个参数上创建一个本地。
然后，由于包含三个子模块，我们需要三次调用，并将其包装到另一个辅助函数中。
现在，我们准备实现训练循环。在初始化模型参数后，我们创建和。分布式优化器将获取参数列表，找到所有不同的所有者工作节点，并使用给定参数（即，在本例中为）在每个所有者工作节点上创建给定的本地优化器（即，您也可以使用其他本地优化器）。
在训练循环中，首先创建一个分布式自动求导上下文，这将帮助分布式自动求导引擎找到梯度和涉及的 RPC 发送/接收函数。分布式自动求导引擎的设计细节可以在其设计说明中找到。然后，启动前向传播，就像是一个本地模型，然后运行分布式反向传播。对于分布式反向传播，您只需要指定一个根列表，在本例中，它是损失。分布式自动求导引擎将自动遍历分布式图并正确写入梯度。接下来，在分布式优化器上运行函数，这将联系到所有涉及的本地优化器来更新模型参数。与本地训练相比，一个小的区别是您不需要运行，因为每个自动求导上下文都有专用空间来存储梯度，并且由于我们每次迭代创建一个上下文，来自不同迭代的梯度不会累积到相同的集合中。
最后，让我们添加一些粘合代码来启动参数服务器和训练器进程。
原文：
译者：飞龙
协议：CC BY-NC-SA 4.0
作者：Rohan Varma
注意
在github中查看并编辑本教程。
先决条件：
- PyTorch 分布式概述
- RPC API 文档
本教程演示了使用 PyTorch 的分布式 RPC 框架实现参数服务器的简单示例。参数服务器框架是一种范式，其中一组服务器存储参数，例如大型嵌入表，几个训练器查询参数服务器以检索最新的参数。这些训练器可以在本地运行训练循环，并偶尔与参数服务器同步以获取最新的参数。要了解更多关于参数服务器方法的信息，请查看这篇论文。
使用分布式 RPC 框架，我们将构建一个示例，其中多个训练器使用 RPC 与同一参数服务器通信，并使用RRef来访问远程参数服务器实例上的状态。每个训练器将通过在多个节点之间的自动求导图上进行分布式反向传递的拼接来启动其专用的反向传递。
注意：本教程涵盖了分布式 RPC 框架的使用，该框架对于将模型分割到多台机器上或实现参数服务器训练策略非常有用，其中网络训练器获取托管在不同机器上的参数。如果您想要在多个 GPU 上复制模型，请参阅分布式数据并行教程。还有另一个RPC 教程，涵盖了强化学习和 RNN 用例。
让我们从熟悉的开始：导入所需的模块并定义一个简单的 ConvNet，该网络将在 MNIST 数据集上进行训练。下面的网络主要采用自pytorch/examples repo中定义的网络。
接下来，让我们定义一些有用的辅助函数，这些函数将对我们脚本的其余部分很有用。以下使用rpc_sync和RRef来定义一个函数，该函数在远程节点上调用给定对象的方法。在下面，我们对远程对象的句柄由参数给出，并在拥有节点上运行它：。在调用节点上，我们通过使用同步运行此命令，这意味着我们将阻塞直到收到响应。
现在，我们准备定义我们的参数服务器。我们将子类化并保存一个句柄到我们上面定义的网络。我们还将保存一个输入设备，这将是在调用模型之前将输入传输到的设备。
接下来，我们将定义我们的前向传递。请注意，无论模型输出的设备如何，我们都将输出移动到 CPU，因为分布式 RPC 框架目前仅支持通过 RPC 发送 CPU 张量。由于调用方/被调用方可能存在不同设备（CPU/GPU），我们故意禁用了通过 RPC 发送 CUDA 张量，但可能会在未来版本中支持。
接下来，我们将定义一些对训练和验证有用的杂项函数。首先，将接收一个分布式自动求导上下文 ID，并调用 API 来检索分布式自动求导计算的梯度。更多信息可以在分布式自动求导文档中找到。请注意，我们还会遍历结果字典，并将每个张量转换为 CPU 张量，因为目前框架只支持通过 RPC 发送张量。接下来，将遍历我们的模型参数，并将它们包装为一个（本地）RRef。这个方法将被训练节点通过 RPC 调用，并返回要优化的参数列表。这是Distributed Optimizer的输入要求，它要求所有必须优化的参数作为列表。
最后，我们将创建方法来初始化我们的参数服务器。请注意，在所有进程中只会有一个参数服务器实例，并且所有训练器将与同一个参数服务器通信并更新相同的存储模型。如中所示，服务器本身不会采取任何独立的行动；它会等待来自训练器（尚未定义）的请求，并通过运行请求的函数来响应它们。
请注意，上面的不会立即关闭参数服务器。相反，它将等待所有工作节点（在这种情况下是训练器）也调用。这样我们就可以保证在所有训练器（尚未定义）完成训练过程之前，参数服务器不会下线。
接下来，我们将定义我们的类。这也将是的子类，我们的方法将使用 API 来获取一个 RRef，或者远程引用，到我们的参数服务器。请注意，这里我们不会将参数服务器复制到我们的本地进程，相反，我们可以将看作是指向在单独进程中运行的参数服务器的分布式共享指针。
接下来，我们将定义一个名为的方法。为了激发对这个方法的需求，值得阅读DistributedOptimizer上的文档，特别是 API 签名。优化器必须传递一个要优化的远程参数的列表，所以这里我们获取必要的。由于给定的与唯一的远程工作节点进行交互，我们只需在上调用。我们使用在类中定义的方法。这个方法将返回一个需要被优化的参数的列表。请注意，在这种情况下，我们的不定义自己的参数；如果定义了，我们还需要将每个参数包装成一个，并将其包含在输入到中。
现在，我们准备定义我们的方法，它将调用（同步）RPC 来运行在上定义的网络的前向传播。请注意，我们传入，这是对我们的的远程句柄，到我们的 RPC 调用中。这个调用将发送一个 RPC 到我们的正在运行的节点上，调用传播，并返回对应于模型输出的。
我们的训练器已经完全定义好了，现在是时候编写我们的神经网络训练循环，该循环将创建我们的网络和优化器，运行一些输入通过网络并计算损失。训练循环看起来很像本地训练程序的循环，但由于我们的网络分布在多台机器上，所以有一些修改。
在下面，我们初始化我们的并构建一个。请注意，如上所述，我们必须传入所有全局（参与分布式训练的所有节点）参数，我们希望进行优化。此外，我们传入要使用的本地优化器，本例中为 SGD。请注意，我们可以像创建本地优化器一样配置底层优化算法 - 所有的参数都将被正确转发。例如，我们传入一个自定义学习率，该学习率将用作所有本地优化器的学习率。
接下来，我们定义我们的主要训练循环。我们循环遍历 PyTorch 的DataLoader提供的可迭代对象。在编写典型的前向/后向/优化器循环之前，我们首先将逻辑包装在Distributed Autograd context中。请注意，这是为了记录模型前向传递中调用的 RPC，以便构建一个适当的图，其中包括在后向传递中参与的所有分布式工作节点。分布式自动求导上下文返回一个，用作累积和优化与特定迭代对应的梯度的标识符。
与调用典型的不同，后者会在本地工作节点上启动后向传递，我们调用并传入我们的以及，这是我们希望从根开始进行后向传递的位置。此外，我们将这个传递给我们的优化器调用，这是必要的，以便能够查找由此特定后向传递计算的相应梯度跨所有节点。
接下来，我们简单地计算模型在训练完成后的准确率，就像传统的本地模型一样。但是，请注意，我们在上面传递给此函数的是的一个实例，因此前向传递以透明方式调用 RPC。
接下来，类似于我们为定义作为主循环的方式，负责初始化 RPC，让我们为训练器定义一个类似的循环。不同之处在于我们的训练器必须运行我们上面定义的训练循环：
请注意，类似于，默认情况下会等待所有工作节点，包括训练器和参数服务器，调用后，该节点才会退出。这确保节点被优雅地终止，而不会在另一个节点期望其在线时离线。
我们现在已经完成了训练器和参数服务器特定的代码，剩下的就是添加代码来启动训练器和参数服务器。首先，我们必须接收适用于我们参数服务器和训练器的各种参数。对应于将参与训练的节点的总数，是所有训练器和参数服务器的总和。我们还必须为每个单独的进程传入一个唯一的，从 0（我们将在其中运行单个参数服务器）到。和是可以用来识别 0 级进程运行位置的参数，并将被各个节点用于发现彼此。为了在本地测试此示例，只需将和相同的传递给所有生成的实例。请注意，出于演示目的，此示例仅支持 0-2 个 GPU，尽管该模式可以扩展以利用更多的 GPU。
现在，我们将根据我们的命令行参数创建相应于参数服务器或训练器的进程。如果我们传入的等级为 0，则将创建一个，否则将创建一个。请注意，我们使用启动一个子进程，对应于我们要执行的函数，并在主线程中使用等待该进程的完成。在初始化我们的训练器时，我们还使用 PyTorch 的dataloaders来指定 MNIST 数据集上的训练和测试数据加载器。
要在本地运行示例，请在单独的终端窗口中为服务器和每个要生成的工作节点运行以下命令：。例如，对于世界大小为 2 的主节点，命令将是。然后可以在单独的窗口中使用命令启动训练器，这将开始使用一个服务器和一个训练器进行训练。请注意，本教程假定训练使用 0 到 2 个 GPU 进行，可以通过将传递到训练脚本中进行配置。
您可以通过命令行参数和传入地址和端口，以指示主工作节点正在侦听的地址和端口，例如，用于测试训练器和主节点在不同机器上运行的功能。
到此这篇argparser怎么安装（arg10.2安装步骤）的文章就介绍到这了,更多相关内容请继续浏览下面的相关推荐文章，希望大家都能在编程的领域有一番成就！版权声明：
本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如若内容造成侵权、违法违规、事实不符，请将相关资料发送至xkadmin@xkablog.com进行投诉反馈，一经查实，立即处理！
转载请注明出处，原文链接：https://www.xkablog.com/qdvuejs/22201.html