使用 Amazon EKS 和 Torch Distributed Elastic 进行分布式训练

由柏拉图重新发布

关注： 0

随着许多行业的数据量不断增长，分布式深度学习模型训练变得越来越重要。现在，计算机视觉和自然语言处理中的许多应用都需要训练深度学习模型，这些模型的复杂性呈指数级增长，并且通常需要使用数百 TB 的数据进行训练。然后，使用庞大的云基础架构来扩展此类大型模型的训练变得很重要。

开发人员可以使用 PyTorch 等开源框架轻松设计直观的模型架构。然而，由于编排复杂性增加，跨多个节点扩展这些模型的训练可能具有挑战性。

分布式模型训练主要包括两种范式：

模型并行 – 在模型并行训练中，模型本身太大，单颗GPU的内存都装不下，需要多颗GPU来训练模型。 Open AI 的 GPT-3 模型具有 175 亿个可训练参数（大小约为 350 GB）就是一个很好的例子。
数据并行 – 在数据并行训练中，模型可以驻留在单个 GPU 中，但由于数据量很大，训练模型可能需要数天或数周的时间。将数据分布在多个 GPU 节点上可以显着减少训练时间。

在这篇文章中，我们提供了一个示例架构来使用 Torch 分布式弹性以分布式数据并行方式使用的框架 Amazon Elastic Kubernetes服务（亚马逊 EKS）。

先决条件

要复制本文中报告的结果，唯一的先决条件是 AWS 账户。在此帐户中，我们创建一个 EKS 集群和一个适用于Lustre的Amazon FSx 文件系统。我们还将容器图像推送到 Amazon Elastic Container注册账户中的 (Amazon ECR) 存储库。在整个帖子中根据需要提供了设置这些组件的说明。

EKS 集群

Amazon EKS 是一种托管容器服务，用于在 AWS 上运行和扩展 Kubernetes 应用程序。借助 Amazon EKS，您可以使用最新的亚马逊弹性计算云 (Amazon EC2) 实例，无需安装、操作和维护您自己的控制平面或节点。这是一个受欢迎的编排用于机器学习 (ML) 和 AI 工作流程。 AWS 中典型的 EKS 集群如下图所示。

我们发布了一个开源项目，适用于 EKS 的 AWS 开发运营 (aws-do-eks)，它提供了大量易于使用和可配置的脚本和工具来配置 EKS 集群和运行分布式训练作业。该项目是按照以下原则建造的做框架: 简单、灵活、通用。您可以使用 eks配置文件文件，然后通过运行 eks-create.sh 脚本。详细说明在 GitHub回购.

使用 Torch 分布式弹性训练 PyTorch 模型

Torch Distributed Elastic (TDE) 是一个本地 PyTorch 库，用于训练大规模深度学习模型，其中基于可用性动态扩展计算资源至关重要。这 Kubernetes 的 TorchElastic 控制器是 TDE 的原生 Kubernetes 实现，可自动管理 TDE 训练所需的 Pod 和服务的生命周期。它允许在训练期间根据需要动态扩展计算资源。它还通过从节点故障中恢复作业来提供容错训练。

在这篇文章中，我们讨论了训练 PyTorch 的步骤 EfficientNet-B7 和残差网络50 模型使用影像网使用 TDE 以分布式方式获取数据。我们使用 PyTorch 分布式数据并行 API 和 Kubernetes TorchElastic 控制器，并在包含多个 GPU 节点的 EKS 集群上运行我们的训练作业。下图显示了此模型训练的架构图。

TorchElastic for Kubernetes 主要由两个组件组成：TorchElastic Kubernetes 控制器（TEC）和参数服务器（etcd）。控制器负责监控和管理训练作业，参数服务器跟踪工作节点以进行分布式同步和对等发现。

为了让训练 pod 访问数据，我们需要一个可以被每个 pod 挂载的共享数据卷。共享卷的一些选项容器存储接口 (CSI) 驱动程序包含在适用于 EKS 的 AWS 开发运营，那恭喜你，亚马逊弹性文件系统（亚马逊 EFS）和 FSx 光泽.

集群设置

在我们的集群配置中，我们将一个 c5.2xlarge 实例用于系统 pod。我们使用三个 p4d.24xlarge 实例作为工作 pod 来训练 EfficientNet 模型。对于 ResNet50 训练，我们使用 p3.8xlarge 实例作为工作 pod。此外，我们使用 FSx 共享文件系统来存储我们的训练数据和模型工件。

AWS p4d.24xlarge 实例配备弹性织物适配器 (EFA) 提供节点之间的网络连接。我们稍后会在帖子中详细讨论 EFA。要通过 EFA 进行通信，我们需要通过 .yaml 文件配置集群设置。一个示例文件在 GitHub 存储库中提供。

正确配置此 .yaml 文件后，我们可以使用 GitHub 存储库中提供的脚本启动集群：

./eks-create.sh

参考 GitHub回购有关详细说明。

在 p4d.24xlarge 和 p3.8xlarge 上运行作业几乎没有区别。本文中描述的步骤适用于两者。唯一的区别是 EFA 在 p4d.24xlarge 实例上的可用性。对于像 ResNet50 这样的小型模型，与 EFA 网络相比，标准网络对训练速度的影响最小。

FSx for Lustre 文件系统

FSx 专为高性能计算工作负载而设计，并使用固态驱动器存储卷提供亚毫秒级延迟。我们选择 FSx 是因为它在我们扩展到大量节点时提供了更好的性能。需要注意的一个重要细节是 FSx 只能存在于单个可用区中。因此，访问 FSx 文件系统的所有节点都应与 FSx 文件系统存在于同一可用区中。实现此目的的一种方法是在创建集群之前在集群 .yaml 文件中为特定节点组指定相关的可用区。或者，我们可以在集群搭建完成后，为这些节点修改 Auto Scaling Group 的网络部分，并将其限制为使用单个子网。这可以在 Amazon EC2 控制台上轻松完成。

假设 EKS 集群已启动并运行，并且可用区的子网 ID 已知，我们可以通过在 fsx配置文件文件中所述自述并运行部署文件脚本中 FSX 文件夹。这将为访问文件系统设置正确的策略和安全组。该脚本还安装 CSI 驱动程序将 FSx 作为守护程序集。最后，我们可以通过应用单个 .yaml 文件在 Kubernetes 中创建 FSx 持久卷声明：

kubectl apply -f fsx-pvc-dynamic.yaml

这会在指定的可用区中创建一个 FSx 文件系统 fsx.conf 文件，并创建一个持久卷声明 fsx-pvc，可以由集群中的任何 pod 以读写多 (RWX) 方式挂载。

在我们的实验中，我们使用了完整的 ImageNet 数据，其中包含超过 12 万张训练图像，分为 1,000 个类别。数据可以从 ImageNet 网站. 原来的 TAR 球有几个目录，但是对于我们的模型训练，我们只对 ILSVRC/Data/CLS-LOC/，其中包括 train 和 val 子目录。在训练之前，我们需要重新排列图像中的图像 val 子目录以匹配 PyTorch 所需的目录结构图片文件夹班级。这可以使用一个简单的 Python脚本在下一步将数据复制到持久卷之后。

从一个复制数据亚马逊简单存储服务 (Amazon S3) 存储桶到 FSx 文件系统，我们创建一个包含此任务脚本的 Docker 映像。示例 Dockerfile 和 shell 脚本包含在沪深 GitHub 存储库中的文件夹。我们可以使用 build.sh 脚本，然后使用 push.sh 脚本。在使用这些脚本之前，我们需要在 .env GitHub 存储库的根文件夹中的文件。将 Docker 映像推送到 Amazon ECR 后，我们可以启动一个 pod，通过应用相关的 .yaml 文件来复制数据：

kubectl apply -f fsx-data-prep-pod.yaml

pod 自动运行脚本数据准备.sh 将数据从 Amazon S3 复制到共享卷。由于 ImageNet 数据有超过 12 万个文件，因此复制过程需要几个小时。 Python 脚本 imagenet_data_prep.py 也运行重新排列 val PyTorch 预期的数据集。

网络加速

我们可以将 Elastic Fabric Adapter (EFA) 与支持的 EC2 实例类型加速集群中 GPU 节点之间的网络流量。这在运行标准网络通信可能成为瓶颈的大型分布式训练作业时非常有用。我们在此处使用的用于在 EKS 集群中部署和测试 EFA 设备插件的脚本包含在 efa 设备插件 GitHub 存储库中的文件夹。要在 EKS 集群中启用 EFA 作业，除了集群节点具有必要的硬件和软件外，还需要将 EFA 设备插件部署到集群，并且您的作业容器需要具有兼容的 CUDA 和 NCCL 版本安装。

为了演示在 p4d.24xlarge 实例上运行 NCCL 测试和评估 EFA 的性能，我们首先必须通过运行相应的部署文件脚本中 mpi 运算符文件夹。然后我们运行部署文件脚本并更新测试-efa-nccl.yaml 表现出对资源的限制和请求 vpc.amazonaws.com 设置为 4。 p4d.24xlarge 节点中的四个可用 EFA 适配器捆绑在一起以提供最大吞吐量。

运行 kubectl apply -f ./test-efa-nccl.yaml 应用测试，然后显示测试 pod 的日志。日志输出中的以下行确认正在使用 EFA：

NCCL INFO NET/OFI Selected Provider is efa

测试结果应类似于以下输出：

[1,0]<stdout>:#                                                       out-of-place                       in-place
[1,0]<stdout>:#       size         count      type   redop     time   algbw   busbw  error     time   algbw   busbw  error
[1,0]<stdout>:#        (B)    (elements)                       (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
[1,0]<stdout>:           8             2     float     sum    629.7    0.00    0.00  2e-07    631.4    0.00    0.00  1e-07
[1,0]<stdout>:          16             4     float     sum    630.5    0.00    0.00  1e-07    628.1    0.00    0.00  1e-07
[1,0]<stdout>:          32             8     float     sum    627.6    0.00    0.00  1e-07    628.2    0.00    0.00  1e-07
[1,0]<stdout>:          64            16     float     sum    633.6    0.00    0.00  1e-07    628.4    0.00    0.00  6e-08
[1,0]<stdout>:         128            32     float     sum    627.5    0.00    0.00  6e-08    632.0    0.00    0.00  6e-08
[1,0]<stdout>:         256            64     float     sum    634.5    0.00    0.00  6e-08    636.5    0.00    0.00  6e-08
[1,0]<stdout>:         512           128     float     sum    634.8    0.00    0.00  6e-08    635.2    0.00    0.00  6e-08
[1,0]<stdout>:        1024           256     float     sum    646.6    0.00    0.00  2e-07    643.6    0.00    0.00  2e-07
[1,0]<stdout>:        2048           512     float     sum    745.0    0.00    0.01  5e-07    746.0    0.00    0.01  5e-07
[1,0]<stdout>:        4096          1024     float     sum    958.2    0.00    0.01  5e-07    955.8    0.00    0.01  5e-07
[1,0]<stdout>:        8192          2048     float     sum    963.0    0.01    0.02  5e-07    954.5    0.01    0.02  5e-07
[1,0]<stdout>:       16384          4096     float     sum    955.0    0.02    0.03  5e-07    955.5    0.02    0.03  5e-07
[1,0]<stdout>:       32768          8192     float     sum    975.5    0.03    0.06  5e-07   1009.0    0.03    0.06  5e-07
[1,0]<stdout>:       65536         16384     float     sum   1353.4    0.05    0.09  5e-07   1343.5    0.05    0.09  5e-07
[1,0]<stdout>:      131072         32768     float     sum   1395.9    0.09    0.18  5e-07   1392.6    0.09    0.18  5e-07
[1,0]<stdout>:      262144         65536     float     sum   1476.7    0.18    0.33  5e-07   1536.3    0.17    0.32  5e-07
[1,0]<stdout>:      524288        131072     float     sum   1560.3    0.34    0.63  5e-07   1568.3    0.33    0.63  5e-07
[1,0]<stdout>:     1048576        262144     float     sum   1599.2    0.66    1.23  5e-07   1595.3    0.66    1.23  5e-07
[1,0]<stdout>:     2097152        524288     float     sum   1671.1    1.25    2.35  5e-07   1672.5    1.25    2.35  5e-07
[1,0]<stdout>:     4194304       1048576     float     sum   1785.1    2.35    4.41  5e-07   1780.3    2.36    4.42  5e-07
[1,0]<stdout>:     8388608       2097152     float     sum   2133.6    3.93    7.37  5e-07   2135.0    3.93    7.37  5e-07
[1,0]<stdout>:    16777216       4194304     float     sum   2650.9    6.33   11.87  5e-07   2649.9    6.33   11.87  5e-07
[1,0]<stdout>:    33554432       8388608     float     sum   3422.0    9.81   18.39  5e-07   3478.7    9.65   18.09  5e-07
[1,0]<stdout>:    67108864      16777216     float     sum   4783.2   14.03   26.31  5e-07   4782.6   14.03   26.31  5e-07
[1,0]<stdout>:   134217728      33554432     float     sum   7216.9   18.60   34.87  5e-07   7240.9   18.54   34.75  5e-07
[1,0]<stdout>:   268435456      67108864     float     sum    12738   21.07   39.51  5e-07    12802   20.97   39.31  5e-07
[1,0]<stdout>:   536870912     134217728     float     sum    24375   22.03   41.30  5e-07    24403   22.00   41.25  5e-07
[1,0]<stdout>:  1073741824     268435456     float     sum    47904   22.41   42.03  5e-07    47893   22.42   42.04  5e-07
[1,4]<stdout>:test-efa-nccl-worker-0:33:33 [4] NCCL INFO comm 0x7fd4a0000f60 rank 4 nranks 16 cudaDev 4 busId 901c0 - Destroy COMPLETE
[1,0]<stdout>:# Out of bounds values : 0 OK
[1,0]<stdout>:# Avg bus bandwidth    : 8.23785

我们可以在测试结果中观察到最大吞吐量约为 42 GB/秒，平均总线带宽约为 8 GB。

我们还在启用单个 EFA 适配器以及不启用 EFA 适配器的情况下进行了实验。所有结果总结在下表中。

EFA 适配器数量	Net/OFI 选择的提供商	平均带宽 (GB/s)	最大限度。带宽 (GB/s)
4	全民教育	8.24	42.04
1	全民教育	3.02	5.89
0	插座	0.97	2.38

我们还发现，对于像 ImageNet 这样的相对较小的模型，使用加速网络可以减少每个 epoch 的训练时间，在 5 的批大小下仅减少 8-64%。对于较大的模型和较小的批大小，当需要增加权重的网络通信时，使用加速联网影响更大。我们观察到使用批量大小为 15 的 EfficientNet-B18 训练的 epoch 训练时间减少了 7-1%。EFA 对训练的实际影响将取决于模型的大小。

GPU监控

在运行训练作业之前，我们还可以设置亚马逊CloudWatch 在训练期间可视化 GPU 利用率的指标。了解资源是否得到最佳使用或潜在地识别训练过程中的资源匮乏和瓶颈可能会有所帮助。

设置 CloudWatch 的相关脚本位于 GPU指标文件夹。首先，我们创建一个 Docker 镜像 amazon-cloudwatch-agent 和 nvidia-smi. 我们可以在 Dockerfile 中使用 gpu-metrics 文件夹来创建这个图像。假设 ECR 注册表已在 .env 上一步中的文件，我们可以使用构建和推送图像 build.sh 和 push.sh. 在此之后，运行 deploy.sh 脚本自动完成设置。它启动一个守护进程 amazon-cloudwatch-agent 并将各种指标推送到 CloudWatch。 GPU 指标显示在 CWAgent CloudWatch 控制台上的命名空间。其余的集群指标显示在 ContainerInsights 命名空间。

模型训练

PyTorch 培训所需的所有脚本都位于弹性工作 GitHub 存储库中的文件夹。在启动训练作业之前，我们需要运行 etcd 服务器，TEC 将其用于工作人员发现和参数交换。这部署文件脚本中 elasticjob 文件夹正是这样做的。

要在 p4d.24xlarge 实例中利用 EFA，我们需要使用 Amazon ECR 公共图库支持通过 EFA 进行的 NCCL 通信。我们只需要将我们的训练代码复制到这个 Docker 镜像中。这 Dockerfile 在下面样本文件夹创建在 p4d 实例上运行训练作业时使用的图像。和往常一样，我们可以使用编译文件和推.sh 文件夹中的脚本来构建和推送图像。

imagenet-efa.yaml 文件描述了训练作业。此 .yaml 文件设置运行训练作业所需的资源，并使用上一节中设置的训练数据安装持久卷。

这里有几点值得指出。副本数应设置为集群中可用的节点数。在我们的例子中，我们将其设置为 3，因为我们有三个 p4d.24xlarge 节点。在里面 imagenet-efa.yaml 文件， nvidia.com/gpu 资源下的参数和 nproc_per_node 下 args 应该设置为每个节点的 GPU 数量，在 p4d.24xlarge 的情况下为 8。此外，Python 脚本的 worker 参数设置每个进程的 CPU 数量。我们选择它为 4，因为在我们的实验中，这在 p4d.24xlarge 实例上运行时提供了最佳性能。为了最大限度地利用集群中所有可用的硬件资源，这些设置是必要的。

当作业运行时，我们可以在 CloudWatch 中观察集群中所有 GPU 的 GPU 使用情况。以下是我们在集群中使用三个 p4d.24xlarge 节点的训练作业之一的示例。在这里，我们从每个节点中选择了一个 GPU。使用前面提到的设置，在 epoch 的训练阶段，集群中所有节点的 GPU 使用率接近 100%。

为了使用 p50xlarge 实例训练 ResNet3.8 模型，我们需要与使用 p4d.24xlarge 训练 EfficientNet 所描述的完全相同的步骤。我们也可以使用相同的 Docker 镜像。如前所述，p3.8xlarge 实例未配备 EFA。然而，对于 ResNet50 模型，这并不是一个明显的缺点。这 imagenet-fsx.yaml GitHub 存储库中提供的脚本使用 p3.8xlarge 节点类型的适当资源设置训练作业。该作业使用来自 FSx 文件系统的相同数据集。

GPU 缩放

我们进行了一些实验，以观察 EfficientNet-B7 模型的训练时间如何通过增加 GPU 的数量来扩展。为此，我们在每次训练运行的训练 .yaml 文件中将副本数从 1 更改为 3。在使用完整的 ImageNet 数据集时，我们只观察了单个 epoch 的时间。下图显示了我们的 GPU 缩放实验的结果。红色虚线表示如何通过增加 GPU 的数量从使用 8 个 GPU 的运行中减少训练时间。正如我们所看到的，缩放非常接近预期。

同样，我们获得了在 p50xlarge 实例上进行 ResNet3.8 训练的 GPU 缩放图。对于这种情况，我们将 .yaml 文件中的副本从 1 更改为 4。此实验的结果如下图所示。

清理

为了避免与运行空闲实例相关的成本，在模型训练后减少资源非常重要。对于每个创建资源的脚本， GitHub回购提供了一个匹配的脚本来删除它们。要清理我们的设置，我们必须在删除集群之前删除 FSx 文件系统，因为它与集群 VPC 中的子网相关联。要删除 FSx 文件系统，我们只需要运行以下命令（从内部 FSX 文件夹）：

kubectl delete -f fsx-pvc-dynamic.yaml
./delete.sh

请注意，这不仅会删除持久卷，还会删除 FSx 文件系统，并且文件系统上的所有数据都将丢失。当这一步完成后，我们可以使用以下脚本删除集群前文件夹：

./eks-delete.sh

这将删除所有现有的 Pod，移除集群，并删除一开始创建的 VPC。

结论

在这篇文章中，我们详细介绍了在 EKS 集群上运行 PyTorch 分布式数据并行模型训练所需的步骤。这项任务看似艰巨，但适用于 EKS 的 AWS 开发运营由 AWS 的 ML Frameworks 团队创建的项目提供了所有必要的脚本和工具来简化流程并使分布式模型训练易于访问。

有关本文中使用的技术的更多信息，请访问亚马逊EKS 和 Torch 分布式弹性. 我们鼓励您将此处描述的方法应用于您自己的分布式训练用例。

资源

关于作者

伊姆兰·尤努斯 是 AWS ML 框架团队的首席解决方案架构师。他专注于跨 AWS 服务（如 Amazon EKS 和 AWS ParallelCluster）的大规模机器学习和深度学习工作负载。他在深度学习在计算机视觉和工业物联网中的应用方面拥有丰富的经验。 Imran 获得了高能粒子物理学博士学位，在那里他参与了 PB 级实验数据的分析。

亚历克斯·扬库尔斯基 是一位全栈软件和基础架构架构师，喜欢做深入的动手工作。他目前是 AWS 自我管理机器学习的首席解决方案架构师。在他的职位上，他专注于帮助客户在容器驱动的 AWS 服务上实现 ML 和 AI 工作负载的容器化和编排。他也是开源的作者做框架还有一位喜欢应用容器技术来加快创新步伐，同时解决世界上最大挑战的 Docker 船长。在过去的 10 年里，Alex 一直致力于应对气候变化、使 AI 和 ML 民主化、使旅行更安全、医疗保健更好、能源更智能。