当前，我们可以使用多种技术通过少量数据训练出深度学习模型，具体包括针对图像分类任务的迁移学习、少样本学习甚至是一次性学习等，也可以基于预训练的 BERT 或 GPT2 模型对语言模型进行微调。但是，在部分应用用例中我们仍然需要引入大量训练数据。例如，如果当前图像与 ImageNet 数据集内的图像完全不同，或者当前语言语料库只针对特定领域、而非通用类型，那么单凭迁移学习将很难带来理想的模型性能。作为深度学习研究人员，您可能需要从零开始尝试新的思路或方法。在这种情况下，我们必须使用大型数据集训练出大型深度学习模型；在找不到最佳训练方法的情况下，整个过程可能需要几天、几周甚至是几个月。
在本文中，我们将一同了解如何在 Amazon SageMaker 的单一实例之上运行多 GPU 训练，并讨论如何在 Amazon SageMaker 上实现高效多 GPU 与多节点分布式训练。
Horovod 基础知识
在使用大量数据进行模型训练时，最好是将训练作业分配给多个 GPU（单一实例或者多个实例）。深度学习框架提供内置方法以支持多 GPU 训练或分布式训练。但除此之外，还有另外一种实现方法，即直接使用分布式深度学习框架（例如Horovod）。Horovod 是 Uber 公司打造的分布式深度学习开源框架，能够与 TensorFlow、Keras、PyTorch 以及 Apache MXNet 等一线热门深度学习工具包协同使用。Horovod 使用 all-reduce 算法取代以往的参数服务器方法进行快速分布式训练，其中还提供多种优化方法以进一步加快分布式训练的执行速度。关于更多详细信息，请参阅遇见Horovod：面向TensorFlow的Uber开源分布式深度学习框架。
为 Horovod 准备数据
在使用 Horovod 执行训练作业时，Horovod 会为其集群当中的每个 GPU 上的工作节点启动独立的进程（每个 GPU 对应一个工作节点）。例如，如果您使用一个包含 4 GPU 的训练实例（一台 Amazon SageMaker ml.p3.8xlarge 或 Amazon Elastic Compute Cloud (Amazon EC2) p3.8xlarge 实例）运行 Horovod 训练作业，则将对应启动 4 个工作进程。数据集本体已经出于数据并行性的需求而被拆分为多个分片，所有这 4 个工作节点都将分别读取自己的数据集分片。如果有 40000 个训练样本，则每个工作节点将获得 10000 个互不重复的训练样本。如果您使用 Horovod 进行分布式训练甚至是多 GPU 训练，则应事先做好数据分片准备，并指引工作节点从文件系统中读取各个分片。（某些深度学习框架可以自动执行此操作，例如 PyTorch 的 DataParallel 与 DistributedDataParallel）
下图所示，为进行分片存储的两种可行架构。
您可以通过多种不同方式为 Amazon SageMaker 训练作业提供数据集。一种最典型的方法就是将所有数据集存储在 Amazon Simple Storage Service (Amazon S3)存储桶内，并在需要时进行访问。大家当然可以使用共享文件系统（例如 Amazon FSx for Lustre 或 Amazon Elastic File System ，简称 Amazon EFS)实现数据存储，但通过 Amazon SageMaker 内置的两种输入模式（文件模式与管道模式）直接从 Amazon S3 中检索数据能够避免系统产生额外的服务成本。
在文件模式下，当 Amazon SageMaker 启动训练作业后，数据集将从指定的 S3 存储桶被传送至训练实例当中，并将其放置在某个特定目录之内。但如果您使用的数据集极为庞大，那么将对象从存储桶复制至训练实例的存储上往往需要耗费很长时间，而且直到数据传输完成，您的训练作业才会真正开始。这会在某些情况下拖慢机器学习（ML）的执行流程，甚至影响到创新或研究工作的项目进度。
另外，大家也可以通过管道模式直接访问存储在 Amazon S3 中的数据集。管道模式在训练实例与 S3 存储桶之间创建直接输入管道，并允许训练进程直接访问对象，这就消除了在训练开始之前将所有对象复制至训练实例中的工作。要对给定 Amazon S3 URI 中的数据集以管道模式加以访问，请在创建 Amazon SageMaker Estimator 时将输入模式设置为 Pipe ，具体参见以下代码：
在管道模式下，训练数据将作为 FIFO 流的形式进行交付。TensorFlow 扩展的 dataset 类极大降低了访问流数据集的难度。关于管道模式与 TensorFlow 的更多详细信息，请参阅在Amazon SageMaker上使用高速管道模式加快模型训练，以及Amazon SageMaker TensorFlow扩展 GitHub repo。
配合 Horovod 使用管道模式
当您配合 Horovod 使用管道模式执行单机多卡或者多机多卡的分布式训练时，有一点需要特别注意。下图所示为这类场景的基本架构。
管道模式将数据从 Amazon S3 流式的传送到训练实例当中的 Unix 命名管道/FIFOs 当中。一个 FIFO 文件仅支持一对写入/读取程序，且每轮训练周期内我们只能为一条通道创建一个 FIFO 文件。通常，人们会为训练数据集定义一条通道，并为验证或测试数据集定义另一条单独的通道，而后将这些输入通道作为 Amazon SageMaker Estimator 中 fit() 函数的参数传递至训练作业。详见以下代码：
这种方式在 Horovod 多 GPU 训练场景下又会造成怎样的影响？简而言之，使用 Horovod 在多 GPU 训练作业中启动的各个进程，会相互争用单一 FIFO，导致多个进程无法同时访问这些 FIFO。而且由于同一时间内只有单一工作进程能够访问 FIFO，且在完成训练作业之前不会释放句柄，就导致所有其他工作进程无法从该 FIFO 中读取数据，最终令训练作业陷入死锁式的无限循环。如果您看到类似于以下形式的重复提示消息，则表明您遇到了这样的问题：
您可以对 S3 存储桶中的数据集进行分片，且数量与用于训练作业的 GPU 数量相对应。如果您拥有 4000 个 TensorFlow 记录文件，且使用一台带有 4 GPU 的 ml.p3.8xlarge 实例进行模型训练，则可以为互不重复的 1000 个 TensorFLow 记录文件设定不同的前缀，如以下代码所示：
使用SharedByS3Key 作为 Amazon S3 数据类型分配方式进行的数据集分片方法，并不完全适用于 Horovod。 这是因为在使用SharedByS3Key 时，分片只会以实例为单位、而非以工作进程为单位进行，且实例中的工作进程与 GPU 的数量保持一致。同样的，各个实例仍然只拥有一条输入通道。因此，大家需要将数据集的分片数量，设定为与 Horovod 集群内 GPU 数相同。
接下来，我们需要为 Amazon SageMaker 训练定义四条输入通道，具体参见以下代码：
ShuffleConfig 将确保根据每个训练轮次，对 Amazon S3 前缀下各文件的使用顺序进行随机分配。关于更多详细信息，请参阅ShuffleConfig。
在 Amazon SageMaker Estimator 上调用fit方法时，请使用以下通道定义：
对于验证及测试类任务，我们只需在单一工作进程上运行（通常使用主工作进程或 Rank 0 工作进程）。在这里，我们不需要设置多条验证或测试通道。但如果您使用 tf.keras.model.fit() 函数进行训练，则训练会在只有一个 Horovod 工作进程进行验证时停止（关于更多详细信息，请参阅 Horovod GitHub repo 上的 issue #600）。如果需要使用 tf.keras.model.fit() 进行验证，大家还应为各验证数据集提供对应的输入通道（类似于训练输入通道）。请注意，截至 2020 年 7 月，管道模式下训练作业的输入通道总数上限为 20 个。具体请参见以下代码：
相较于直接使用 S3 存储桶前缀，我们在这里可以使用包含有对象键列表的普通ManifestFile。关于更多详细信息，请参阅输入数据。
在训练代码中使用数据通道
在训练脚本中，我们需要强制要求各个 Horovod 工作进程只访问属于它自己的数据集分片，确保两个工作进程不会访问同一输入通道。在本文的用例中，我们将使用从 0 开始的索引定义各输入通道名称。为此，我们可以使用hvd.rank() 函数，由其为当前工作进程在集群范围之内提供唯一的排名索引，且排名同样从 0 开始（请参考以下代码中的第 13 行）。在本文示例中，我们使用 Amazon SageMaker TensorFlow 扩展 PipeModeDataset。对于其他深度学习框架，请在每个训练轮次中从名为 /opt/ml/input/data/[channel_name]_${epoch} 的 FIFO 文件中读取数据。关于更多示例，请参见 GitHub repo。
在包含一个或多个实例的 Horovod 集群中，排名分配方式为从 0 开始，至 GPU 数量-1 结束。只要正确定义了输入通道的名称并从 0 开始使用索引，我们就不必分神管理各实例或名位的排列顺序。
使用 Tensorboard 进行监控
在对训练进程加以灵活监控方面，我们可以在每个训练轮次结束时首先将日志上传至 S3 存储桶，再通过任意远程计算实例调用 Tensorboard。为此，我们需要创建一项回调以将本地日志推送至 S3 存储桶路径，此路径仅限于运行在 Horovod 上的主（Rank 0）计算节点。具体请参见以下代码：
通过将训练日志存储在 S3 存储桶内，大家可以在任意服务器上运行 Tensorboard，包括 EC2 实例、Amazon SgaeMaker notebook 实例甚至是本地计算机，并为该 Tensorboard 托管服务器提供访问 Amazon S3 日志对象的权限。为了支持从 Amazon S3 源处直接提取日志数据，您的 Tensorboard 必须为 1.14.0 或者更高版本。以下命令行使用的是位于us-east-1区域内 S3 存储桶上的日志记录：
如果您是在 Amazon SageMaker notebook 实例上运行以上命令，则可通过 https://<SageMaker-notebook-instance-name>.notebook.<notebook-region>.sagemaker.aws/proxy/6006/ 完成访问。
资源清理
在完成本文中分布式训练作业之后，请清理相应资源以避免其后续产生额外费用，包括 S3 存储桶、FSx for Lustre 以及各个 Amazon SageMaker 实例。
总结
在 Amazon SageMaker 上以管道模式使用 Horovod 的多 GPU 或分布式训练方法，能够为数据集的各个分片创建独立的训练通道并在数据通道内访问对应分片，借此实现大规模模型训练。这种方式能够缩短在实际训练开始之前将数据集传输至训练实例所占用的时间，因此特别适用于具有大规模训练数据集的 Amazon SageMaker 训练场景。
关于在 Amazon SageMaker 上运行的完整训练示例（管道模式加 Horovod），请参阅GitHub repo。
作者介绍：
Amazon 机器学习解决方案实验室数据科学家。他运用机器学习与深度学习技术帮助客户解决各类业务问题，同时帮助他们提升相关技能水平。
Amazon 机器学习解决方案实验室深度学习架构师。凭借在 AWS 上为全球企业设计工作负载的丰富经验，他目前专门负责为客户的新型业务问题设计并实施机器学习解决方案。
Amazon 机器学习解决方案实验室数据科学家。他与各行各业的客户开展合作，设计并构建起能够产生实际商业价值的自动化算法模型。
本文转载自亚马逊 AWS 官方博客。
原文链接：
在 Amazon SageMaker 管道模式下使用 Horovod 实现多 GPU 分布式训练
分布式云行业实践指南（2023）
业内首个分布式云行业实践方法论、工具箱。
评论