使用 Amazon SageMaker 异构集群提高模型训练的性价比

这篇文章是与 Mobileye 的 Chaim Rand 共同撰写的。

某些机器学习 (ML) 工作负载（例如训练计算机视觉模型或强化学习）通常涉及将神经网络模型训练的 GPU 或加速器密集型任务与数据预处理（例如图像增强）的 CPU 密集型任务相结合。当两种类型的任务在同一实例类型上运行时，数据预处理会在 CPU 上遇到瓶颈，从而导致 GPU 利用率降低。随着时间的推移，这个问题变得更加严重，因为新一代 GPU 的吞吐量增长速度比 CPU 更快。

为了解决这个问题，我们在 2022 年 XNUMX 月 推出 异构集群 亚马逊SageMaker 模型训练，使您能够启动在单个作业中使用不同实例类型的训练作业。这允许将数据预处理管道的一部分卸载到 计算优化 实例类型，而深度神经网络 (DNN) 任务继续在 GPU或加速计算 实例类型。 我们的基准测试显示，在受 CPU 限制的 TensorFlow 计算机视觉模型训练中启用异构集群后，性价比优势高达 46%。

对于类似的用例， Mobileye在一家自动驾驶汽车技术开发公司，分享如下：

“通过将受 CPU 限制的深度学习计算机视觉模型训练转移到在多个实例类型（CPU 和 GPU/ML 加速器）上运行，使用 tf.data.service 基于我们构建的解决方案，我们成功地将训练时间减少了 40%，同时将训练成本降低了 30%。我们很高兴异构集群允许我们在 Amazon SageMaker 上运行此解决方案。”

— 人工智能工程，Mobileye

在这篇文章中，我们将讨论以下主题：

• 异构集群如何帮助消除 CPU 瓶颈
• 何时使用异构集群以及其他替代方案
• PyTorch 和 TensorFlow 中的参考实现
• 性能基准测试结果
• Mobileye 的异构集群

AWS的 加速计算实例 系列包括 AWS 定制芯片的加速器（AWS 推理, AWS 培训), 英伟达 (图形处理器）， 高迪加速器 来自 Habana Labs（英特尔公司）。请注意，在本文中，我们交替使用术语 GPU 和加速器。

异构集群如何消除数据处理瓶颈

训练深度学习模型的数据科学家旨在最大限度地提高训练成本效率并最大限度地缩短训练时间。为了实现这一目标，一个基本的优化目标是拥有高 GPU 利用率，这是业界最昂贵和最稀缺的资源。 亚马逊弹性计算云 （亚马逊 EC2）实例。对于将经典 GPU 密集型神经网络模型的前向和反向传播与 CPU 密集型任务（例如计算机视觉中的数据处理和增强或在强化学习中运行环境模拟）相结合的 ML 工作负载来说，这可能更具挑战性。这些工作负载最终可能会受到 CPU 的限制，因为现有的加速器部分闲置，因此拥有更多的 CPU 会带来更高的吞吐量以及更快、更便宜的训练。在某些情况下，可以通过切换到具有更高 CPU:GPU 比率的其他实例类型来解决 CPU 瓶颈。但是，在某些情况下，由于实例系列的架构、存储或网络依赖性，可能无法切换到其他实例类型。

在这种情况下，您必须通过混合实例类型来增加 CPU 算力：具有 GPU 的实例与 CPU 实例。总而言之，这会导致总体上更高的 CPU:GPU 比率。直到最近，SageMaker 训练作业还仅限于拥有单一选定实例类型的实例。借助 SageMaker 异构集群，数据科学家可以轻松运行具有多种实例类型的训练作业，从而将部分现有 CPU 任务从 GPU 实例卸载到专用的计算优化 CPU 实例，从而提高 GPU 利用率，并实现更快且成本更低的效果。高效的培训。此外，借助额外的 CPU 能力，您可以将传统上作为训练初步步骤离线完成的预处理任务变成训练工作的一部分。这使得数据预处理和 DNN 训练假设和超参数的迭代和实验速度更快。

例如，考虑强大的 GPU 实例类型 ml.p4d.24xlarge（96 个 vCPU、8 个 NVIDIA A100 GPU），CPU:GPU 比例为 12:1。假设您的训练作业需要 20 个 vCPU 来预处理足够的数据，以保持一个 GPU 100% 的利用率。因此，为了保持所有 8 个 GPU 100% 得到利用，您需要 160 个 vCPU 实例类型。但是，ml.p4d.24xlarge 缺少 64 个 vCPU，即 40%，将 GPU 利用率限制为 60%，如下图左侧所示。添加另一个 ml.p4d.24xlarge 实例会有帮助吗？不会，因为作业的 CPU:GPU 比率将保持不变。

对于异构集群，我们可以添加两个ml.c5.18xlarge（72 vCPU），如图右侧所示。该集群中的净 vCPU 总数为 210 (96+2*72)，导致 CPU:GPU 比率为 30:1。这些计算优化实例中的每一个都将卸载数据预处理 CPU 密集型任务，并允许高效地利用 GPU。尽管 ml.c5.18xlarge 会产生额外成本，但更高的 GPU 利用率可实现更快的处理速度，从而带来更高的性价比优势。

何时使用异构集群以及其他替代方案

在本节中，我们将解释如何识别 CPU 瓶颈，并讨论使用实例类型纵向扩展与异构集群来解决该瓶颈。

识别 CPU 瓶颈的快速方法是监控 CPU 和 GPU 利用率指标 用于 SageMaker 培训工作 亚马逊CloudWatch。您可以从以下位置访问这些视图 AWS管理控制台 在训练作业页面的实例指标超链接中。选择相关指标并将分辨率从 5 分钟切换为 1 分钟。请注意，每个 vCPU 或 GPU 的规模为 100%，因此具有 4 个 vCPU/GPU 的实例的利用率可能高达 400%。下图是 CloudWatch 指标中的一个此类示例，其中 CPU 利用率约为 100%，表明存在 CPU 瓶颈，而 GPU 利用率不足。

要进行详细诊断，请使用以下命令运行训练作业 Amazon SageMaker调试器 通过在使用 SageMaker Python SDK 构建 SageMaker 估算器时添加探查器配置来分析资源利用率状态、统计数据和框架操作。提交训练作业后，查看结果 分析报告 针对CPU瓶颈。

如果您认为较高的 CPU:GPU 计算比可以使您的作业受益，请首先考虑扩展到同一实例系列中的另一种实例类型（如果有）。例如，如果您在 ml.g5.8xlarge（32 个 vCPU、1 个 GPU）上训练模型，请考虑扩展到 ml.g5.16xlarge（64 个 vCPU、1 个 GPU）。或者，如果您使用多 GPU 实例 ml.g5.12xlarge（48 个 vCPU、4 个 GPU）训练模型，请考虑扩展到 ml.g5.24xlarge（96 个 vCPU、4 个 GPU）。请参阅 G5 实例系列规格了解更多详细信息。

有时，扩展并不是一种选择，因为在同一实例系列中没有具有更高 vCPU:GPU 比率的实例类型。例如，如果您在 ml.trn1.32xlarge、ml.p4d.24xlarge 或 ml.g5.48xlarge，您应该考虑用于 SageMaker 模型训练的异构集群。

除了扩展之外，我们还想指出，异构集群还有其他替代方案，例如 NVIDIA DALI，它将图像预处理任务转移到 GPU。欲了解更多信息，请参阅 使用 TensorFlow Data Service、NVIDIA DALI 和其他方法克服数据预处理瓶颈.

为了简化决策，请参考以下流程图。

如何使用 SageMaker 异构集群

要快速开始，您可以直接跳转到本文中提供的 TensorFlow 或 PyTorch 示例。

在本部分中，我们将通过一个简单的示例引导您了解如何使用 SageMaker 异构集群。我们假设您已经知道如何使用 SageMaker Python SDK 和 Estimator 类训练模型。如果没有，请参考 使用 SageMaker Python SDK 继续之前。

在此功能之前，您使用以下命令初始化了训练作业的 Estimator 类 InstanceCount 和 InstanceType 参数，隐式假设您只有一个实例类型（同构集群）。随着异构集群的发布，我们引入了新的 sagemaker.instance_group.InstanceGroup 班级。这表示一组特定实例类型的一个或多个实例，旨在承载逻辑角色（例如数据处理或神经网络优化）。您可以拥有两个或多个组，并为每个实例组指定自定义名称，实例类型，以及每个实例组的实例数量，更多信息请参见 使用 SageMaker Python SDK 和 使用低级 SageMaker API.

定义实例组后，您需要修改训练脚本以读取 SageMaker 训练环境信息 其中包括异构集群配置。该配置包含当前实例组、每个组中当前主机、当前主机所在组及其排名等信息。您可以在训练脚本中构建逻辑，以将实例组分配给某些训练和数据处理任务。此外，您的训练脚本需要处理实例间组通信或分布式数据加载机制（例如， tf.数据.服务 在 TensorFlow 或通用中 gRPC 客户端-服务器）或任何其他框架（例如， Apache Spark).

让我们看一个启动异构训练作业并在运行时读取环境配置的简单示例。

1. 定义和启动训练作业时，我们配置两个实例组，用作 SageMaker 估计器的参数：

   from sagemaker.instance_group import InstanceGroup
   data_group = InstanceGroup("data_group", "ml.c5.18xlarge", 2)
   dnn_group = InstanceGroup("dnn_group", "ml.p4d.24xlarge", 1)
   
   from sagemaker.pytorch import PyTorch
   estimator = PyTorch(...,
       entry_point='launcher.py',
       instance_groups=[data_group, dnn_group]
   )

2. 在入口点训练脚本上（名为 launcher.py），我们读取异构集群配置以确定实例是否运行预处理或 DNN 代码：

   from sagemaker_training import environment
   env = environment.Environment()
   if env.current_instance_group == 'data_group': ...;

至此，我们总结一下 SageMaker 代表您执行的任务以及您负责的任务。

SageMaker 执行以下任务：

1. 根据实例组定义配置不同的实例类型。
2. 在所有或特定实例组上配置输入通道。
3. 将训练脚本和依赖项分发到实例。
4. 在特定实例组（如果已定义）上设置 MPI 集群。

您负责以下任务：

1. 修改您的启动训练作业脚本以指定实例组。
2. 实现分布式数据管道（例如， tf.data.service).
3. 修改您的入口点脚本（请参阅 launcher.py 在示例笔记本中）作为单个入口点，将在所有实例上运行，检测它在哪个实例组中运行，并触发相关行为（例如数据处理或 DNN 优化）。
4. 训练循环结束后，您必须确保入口点进程在所有实例组的所有实例上退出。这很重要，因为 SageMaker 会等待所有实例完成处理，然后才会将作业标记为完成并停止计费。这 launcher.py TensorFlow 和 PyTorch 示例笔记本中的脚本提供了信号数据组实例的参考实现，以便在 DNN 组实例完成其工作时退出。

SageMaker 异构集群的示例笔记本

在本节中，我们总结了 示例笔记本 适用于 TensorFlow 和 PyTorch ML 框架。在笔记本中，您可以找到实现细节、代码工作原理的演练、可在训练脚本中重用的代码片段、流程图和成本比较分析。

请注意，在这两个示例中，您不应期望模型以有意义的方式收敛。我们的目的只是测量以纪元/步时间表示的数据管道和神经网络优化吞吐量。您必须使用自己的模型和数据集进行基准测试，才能产生与您的工作负载相匹配的性价比优势。

使用基于 tf.data.service 的分布式数据加载器 (TensorFlow) 的异构集群

本篇 笔记本 演示如何使用 TensorFlow 实现用于 SageMaker 训练的异构集群 tf.data.service 基于分布式数据管道。我们训练深度学习计算机视觉模型 资源网50 这需要 CPU 密集型数据增强。它用 霍尔沃德 用于多 GPU 分布式数据并行。

我们以两种配置运行工作负载：首先作为同质集群、单个 ml.p4d.24xlarge 实例，使用标准 tf.data 显示导致 GPU 利用率降低的 CPU 瓶颈的管道。在第二次运行中，我们使用 SageMaker 异构集群从单一实例类型切换到两个实例组。此运行将一些数据处理卸载到其他 CPU 实例（使用 tf.data.service).

然后，我们比较同构和异构配置并找到关键的性价比优势。如下表所示，异构作业（86 毫秒/步）的训练速度是同质作业（2.2 毫秒/步）的 192 倍，训练模型的成本降低了 46%。

示例1（TF）	毫升.p4d.24xl	毫升.c5.18xl	每小时价格*	平均步数时间	每步成本	性价比提升
同质	1	0	$37.688	192毫秒	$0.201	.
异质	1	2	$45.032	86毫秒	$0.108	46%

* 每小时价格基于 us-east-1 SageMaker 按需定价

通过利用数据组提供的额外 vCPU 和更快的预处理，可以实现这种加速。请参阅 笔记本 了解更多详细信息和图表。

使用基于 gRPC 客户端-服务器的分布式数据加载器 (PyTorch) 的异构集群

本篇 笔记本 演示了使用异构集群进行 S​​ageMaker 训练的示例工作负载，并使用基于 gRPC 客户端-服务器的分布式数据加载器。此示例使用单个 GPU。我们使用基于以下内容的 PyTorch 模型 官方 MNIST 示例。训练代码已被修改为大量数据预处理。我们在同构和异构集群模式下训练该模型，并比较性价比。

在此示例中，我们假设工作负载无法从多个 GPU 中受益，并且依赖于特定的 GPU 架构（NVIDIA V100）。我们运行同质和异构训练作业，并发现了关键的性价比优势，如下表所示。异构作业（1.19 秒/步）的训练速度比同类作业（6.5 秒/步）快 0.18 倍，训练模型的成本降低了 77%。

示例2（PT）	毫升.p3.2xl	毫升.c5.9xl	每小时价格*	平均步数时间	每步成本	性价比提升
同质	1	0	$3.825	1193毫秒	$0.127	.
异质	1	1	$5.661	184毫秒	$0.029	77%

* 每小时价格基于 us-east-1 SageMaker 按需定价

这是可能的，因为 CPU 数量较多，我们可以使用 32 个数据加载器工作线程（ml.p8xlarge 为 3.2 个）来预处理数据，并使 GPU 经常保持接近 100% 的利用率。请参阅 笔记本 了解更多详细信息和图表。

Mobileye 的异构集群

Mobileye 是英特尔旗下公司，开发高级驾驶辅助系统 (ADAS) 和自动驾驶汽车技术，其目标是彻底改变交通运输行业，使道路更安全并拯救生命。这些技术是通过复杂的计算机视觉 (CV) 模型实现的，这些模型使用 SageMaker 对存储在 亚马逊简单存储服务 （亚马逊 S3）。这些模型使用最先进的深度学习神经网络技术。

我们注意到，对于我们的 CV 模型之一，CPU 瓶颈主要是由于繁重的数据预处理导致 GPU 利用率不足造成的。对于这个特定的工作负载，我们开始寻找替代解决方案，评估基于 EC2 实例的异构集群的分布式数据管道技术，并提出了这两种解决方案的参考实现 TensorFlow 和 PyTorch。 SageMaker 异构集群的发布使我们能够在 SageMaker 上运行此工作负载和类似的工作负载，以实现更高的性价比优势。

需要考虑的事项

随着异构集群功能的推出，SageMaker 在训练作业中混合和匹配实例类型方面提供了更大的灵活性。但是，使用此功能时请考虑以下事项：

• 异构集群功能可通过 SageMaker 使用 PyTorch 和 TensorFlow 框架估计器类。支持的框架为 PyTorch v1.10 或更高版本以及 TensorFlow v2.6 或更高版本。
• 所有实例组共享相同的 Docker 映像。
• 所有实例组共享相同的训练脚本。因此，应修改您的训练脚本以检测它属于哪个实例组并相应地运行 fork。
• 训练实例主机名（例如 alog-1、algo-2 等）是随机分配的，并且不指示它们属于哪个实例组。要获取实例的角色，我们建议在运行时获取其实例组成员身份。这在查看登录时也很重要 云观察，因为日志流名称 [training-job-name]/algo-[instance-number-in-cluster]-[epoch_timestamp] 有主机名。
• 一种分布式训练策略（通常是MPI集群）只能应用于一个实例组。
• SageMaker 托管温水池 和 SageMaker 本地模式 目前不能用于异构集群训练。

结论

在这篇文章中，我们讨论了何时以及如何使用 SageMaker 训练的异构集群功能。我们在实际用例中展示了 46% 的性价比提升，并帮助您快速开始使用分布式数据加载器（tf.data.service 和 gRPC 客户端-服务器）实现。您可以在现有训练脚本中使用这些实现，只需进行最少的代码更改。

首先，请尝试我们的 示例笔记本。要了解有关此功能的更多信息，请参阅 使用异构集群进行训练.

---

关于作者

吉利纳楚姆 是一名高级 AI/ML 专家解决方案架构师，是 EMEA 亚马逊机器学习团队的一员。 Gili 对训练深度学习模型的挑战以及机器学习如何改变我们所知道的世界充满热情。 在业余时间，吉利喜欢打乒乓球。

赫鲁什凯什甘古尔 是 AI/ML 初创公司的首席解决方案架构师，拥有 ML 训练和 AWS 网络方面的专业知识。他帮助自动驾驶汽车、机器人、CV、NLP、MLOps、ML 平台和机器人流程自动化技术领域的初创公司在 AWS 上高效且高效地运营业务。在加入 AWS 之前，Hrushikesh 拥有 20 多年主要围绕云和数据平台的行业经验。

盖尔·奥什里 是 Amazon SageMaker 团队的高级产品经理。 他在机器学习工具、框架和服务方面拥有 7 年的工作经验。

柴兰德 是英特尔公司 Mobileye 的一名机器学习算法开发人员，致力于自动驾驶汽车解决方案的深度学习和计算机视觉技术。看看他的 博客.