使用 SageMaker 进行分布式计算的最佳实践 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

使用 SageMaker 进行分布式计算的最佳实践

本最佳实践页面概要介绍了用于机器学习 (ML) 作业的各种分布式计算。本页面中的分布式计算一词包括用于机器学习任务的分布式训练,以及用于数据处理、数据生成、特征工程和强化学习的并行计算。在本页面中,我们将讨论分布式计算中的常见挑战,以及可用于 SageMaker Training 和 SageMaker Processing 的选项。有关分布式计算的其他阅读材料,请参阅什么是分布式计算?

您可以对 ML 任务进行配置,使之在多个节点(实例)、加速器(NVIDIA GPU、Amazon Trainium 芯片)和 vCPU 核心之间以分布式方式运行。通过运行分布式计算,您可以实现各种目标,例如计算速度更快、处理大型数据集或训练大型 ML 模型。

以下列表涵盖了大规模运行 ML 训练作业时可能面临的常见挑战。

  • 您需要根据 ML 任务、要使用的软件库和计算资源来决定如何分布计算资源。

  • 并非所有 ML 任务都能够轻易地实现分布式处理。此外,并非所有 ML 库都支持分布式计算。

  • 分布式计算并不一定能够实现计算效率的线性提高。尤其是您需要确定,数据 I/O 和 GPU 间通信是否存在瓶颈或导致开销。

  • 分布式计算可能会干扰数值处理并改变模型准确性。特别是对于数据并行神经网络训练,当您纵向扩展到更大的计算集群时需要更改全局批量大小,还需要相应地调整学习率。

SageMaker 提供分布式训练解决方案,以针对各种使用场景来轻松应对此类挑战。请选择以下最适合您使用场景的选项。

选项 1:使用支持分布式训练的 SageMaker 内置算法

SageMaker 提供了内置算法,您可以通过 SageMaker 控制台或 SageMaker Python SDK 直接使用。使用内置算法时,您无需花时间自定义代码、了解模型背后的科学知识,也无需在预置 Amazon EC2 实例上运行 Docker。

SageMaker 内置算法中有一部分支持分布式训练。要检查您选择的算法是否支持分布式训练,请参阅关于内置算法的常用信息表中的可并行化列。一些算法支持多实例分布式训练,而其余的可并行化算法则支持在单个实例中使用多个 GPU 进行并行处理,如可并行化列所示。

选项 2:在 SageMaker 托管的训练或处理环境中运行自定义 ML 代码

SageMaker 作业可以为特定的使用场景和框架实例化分布式训练环境。此环境充当即用型白板,您可以在其中引入自己的 ML 代码并运行。

ML 代码使用深度学习框架时

您可以使用适用于 SageMaker Training 的深度学习容器 (DLC) 来启动分布式训练作业,并通过 SageMaker Python SDK 中的专用 Python 模块,或者通过 SageMaker API 使用 Amazon CLIAmazon SDK for Python (Boto3),以此对分布式训练作业进行编排。SageMaker 为机器学习框架提供训练容器,包括 PyTorchTensorFlowHugging Face TransformersApache MXNet。您可以通过两种方式编写深度学习代码用于分布式训练。

  • SageMaker 分布式训练库

    SageMaker 分布式训练库为神经网络数据并行性和模型并行性提供了 Amazon 托管代码。SageMaker 分布式训练还附带了 SageMaker Python SDK 中内置的启动程序客户端,您无需编写并行启动代码。要了解更多信息,请参阅 SageMaker 的数据并行性库SageMaker 的模型并行性库

  • 开源分布式训练库

    开源框架有自己的分布机制,例如 PyTorch 中的 DistributedDataParallelism (DDP) 或 TensorFlow 中的 tf.distribute 模块。您可以选择在 SageMaker 托管的框架容器中运行这些分布式训练框架。例如,在 SageMaker 中训练 MaskRCNN 示例代码展示了如何在 SageMaker PyTorch 框架容器中使用 PyTorch DDP,以及在 SageMaker TensorFlow 框架容器中使用 Horovod

SageMaker ML 容器还预装了 MPI,因此您可以使用 mpi4py 对入口点脚本进行并行化。如果您启动第三方分布式训练启动程序,或者在 SageMaker 托管训练环境中编写临时并行代码,使用 MPI 集成训练容器是一个很好的选择。

在 GPU 上进行数据并行神经网络训练的注意事项

  • 在适当的时候扩展到多 GPU 和多计算机并行性

    我们经常在多 CPU 或多 GPU 实例上运行神经网络训练作业。每个基于 GPU 的实例通常包含多个 GPU 设备。因此,分布式 GPU 计算可以在具有多个 GPU 的单个 GPU 实例中进行(单节点多 GPU 训练),也可以在每个实例中有多个 GPU 核心的多个 GPU 实例中进行(多节点多 GPU 训练)。单实例训练更易于编写代码和调试,而且节点内 GPU 到 GPU 的吞吐量通常比节点间 GPU 到 GPU 的吞吐量更快。因此,好的做法是先垂直扩展数据并行度(使用带有多个 GPU 的单 GPU 实例),然后在需要时扩展到多个 GPU 实例。这可能不适用于需要大量 CPU(例如,用于数据预处理的大量工作负载)的情况,也不适用于多 GPU 实例中 CPU 与 GPU 比率过低的情况。在所有情况下,您都需要根据自己的 ML 训练需求和工作负载,尝试不同的实例类型组合。

  • 监控收敛质量

    在使用数据并行性训练神经网络时,如果增加 GPU 数量的同时保持每个 GPU 的小批次大小不变,就会导致小批次随机梯度下降 (MSGD) 进程的全局小批次的大小增加。而 MSGD 的小批次大小已知会影响下降噪声和收敛性。为了在保持准确性的同时正确缩放,您需要调整其他超参数,例如学习率 [Goyal 等。(2017)]。

  • 监控 I/O 瓶颈

    在您增加 GPU 的数量时,读取和写入存储的吞吐量也应增加。请确保您的数据来源和管道不会成为瓶颈。

  • 根据需要修改训练脚本

    针对单 GPU 训练编写的训练脚本必须进行修改,才能用于多节点多 GPU 训练。在大多数数据并行性库中,需要修改脚本才能进行以下操作。

    • 向每个 GPU 分配训练数据批次。

    • 使用可以处理多个 GPU 上的梯度计算和参数更新的优化器。

    • 将执行检查点操作的责任分配给特定的主机和 GPU。

ML 代码涉及表格数据处理时

PySpark 是 Apache Spark 的 Python 前端,这是一个开源的分布式计算框架。PySpark 已广泛用于大规模生产工作负载的分布式表格数据处理。如果您要运行表格数据处理代码,请考虑使用 SageMaker Processing PySpark 容器并运行并行作业。您还可以在 Amazon SageMaker Studio(与 Amazon EMRAmazon Glue 集成)中,使用 SageMaker Training 和 SageMaker Processing API 并行运行数据处理作业。

选项 3:编写自己的自定义分布式训练代码

当您向 SageMaker 提交训练或处理作业时,SageMaker Training 和 SageMaker Processing API 会启动 Amazon EC2 计算实例。您可以运行自己的 Docker 容器,也可以在 Amazon 托管容器中安装其他库,以此来自定义实例中的训练和处理环境。有关 Docker 与 SageMaker Training 的更多信息,请参阅调整自己的 Docker 容器以用于 SageMaker使用自己的算法和模型创建容器。有关 Docker 与 SageMaker Processing 的更多信息,请参阅使用自己的处理代码

每个 SageMaker Training 作业环境都包含一个位于 /opt/ml/input/config/resourceconfig.json 的配置文件,每个 SageMaker Processing 作业环境都包含一个位于 /opt/ml/config/resourceconfig.json 的类似配置文件。您的代码可以读取此文件,从而查找 hostnames 和建立节点间通信。要了解更多信息,包括 JSON 文件的架构,请参阅分布式训练配置Amazon SageMaker Processing 如何配置处理容器。您也可以在 SageMaker 中安装和使用第三方分布式计算库,例如 Ray 或 DeepSpeed。

您还可以使用 SageMaker Training 和 SageMaker Processing 运行不需要工作线程间通信的自定义分布式计算。在计算文献中,这些任务通常被描述为易并行不共享。这样的例子包括并行处理数据文件、在不同配置下并行训练模型,或者对记录集合运行批量推理。您可以用 Amazon SageMaker 轻易地将这种不共享使用场景并行化。如果您在有多个节点的集群上启动 SageMaker Training 或 SageMaker Processing 作业,默认情况下,SageMaker 会在所有节点上复制并启动您的训练代码(在 Python 或 Docker 中)。如果任务需要在这样的多个节点上随机分散输入数据,则可在 SageMaker TrainingInput API 的数据输入配置中设置 S3DataDistributionType=ShardedByS3Key 来协调这一任务。

选项 4:并行或按顺序启动多个作业

您还可以将 ML 计算工作流分布到较小的并行或顺序计算任务中,每个任务都通过自己的 SageMaker Training 或 SageMaker Processing 作业表示。对于以下情况或任务,将一个任务拆分为多个作业可能会带来好处:

  • 当每个子任务都有特定的数据通道和元数据条目(例如超参数、模型配置或实例类型)时。

  • 当您在子任务级别实施重试步骤时。

  • 当您在运行工作负载期间改变子任务的配置时,例如增加批次大小来进行训练时。

  • 当您需要运行的 ML 任务用时比单个训练作业允许的最大训练时间(最多 28 天)更长时。

  • 当计算工作流的不同步骤需要不同的实例类型时。

有关超参数搜索的具体案例,请使用 SageMaker 自动模型调优。SageMaker 自动模型调优是一款无服务器参数搜索编排工具,可根据随机、贝叶斯或 HyperBand 搜索逻辑,代表您启动多个训练作业。

此外,要编排多个训练作业,您还可以考虑使用工作流编排工具,例如 SageMaker PipelinesAmazon Step Functions,以及 Amazon Managed Workflows for Apache Airflow (MWAA)SageMaker Workflows 支持的 Apache Airflow。