在 Amazon SageMaker Training 上运行 PyTorch Lightning 和本机 PyTorch DDP，具有 Amazon 搜索功能

这么多的数据，这么少的时间。 机器学习 (ML) 专家、数据科学家、工程师和爱好者在世界各地都遇到了这个问题。 从自然语言处理到计算机视觉，从表格到时间序列，以及介于两者之间的一切，在尽可能多的 GPU 上运行数据时优化速度这一古老的问题激发了无数的解决方案。 今天，我们很高兴地宣布使用原生开源框架的 PyTorch 开发人员的功能，例如 PyTorch闪电 和 PyTorch DDP，这将简化他们通向云的路径。

亚马逊SageMaker SageMaker 模型训练是针对 ML 的完全托管服务，而 SageMaker 模型训练是用于大规模高性能训练的优化计算环境。 SageMaker 模型训练通过无缝控制平面提供远程训练体验，以高性能和低成本轻松训练和重现 ML 模型。 我们很高兴地宣布 SageMaker 培训组合中的新功能使大规模运行 PyTorch 变得更加容易和访问：

1. PyTorch Lightning 现在可以集成到 SageMaker 的分布式数据并行库 只需更改一行代码。
2. SageMaker 模型训练 现在支持带有 NCCL 后端的本地 PyTorch 分布式数据并行，使开发人员能够比以往更轻松地迁移到 SageMaker。

在这篇文章中，我们将讨论这些新功能，并了解 Amazon Search 如何使用 SageMaker 中优化的分布式训练后端运行 PyTorch Lightning，以加快模型训练时间。

在深入亚马逊搜索案例研究之前，对于那些不熟悉的人，我们想提供一些背景知识 SageMaker 的分布式数据并行库. 2020 年，我们开发并推出了用于大规模分布式梯度下降的自定义集群配置，以提高整体集群效率，在 Amazon Science 上介绍为 鲱鱼. SageMaker 分布式数据并行 (SMDDP) 充分利用参数服务器和基于环的拓扑结构，针对 亚马逊弹性计算云 (Amazon EC2) 网络拓扑，包括 全民教育. 对于更大的集群规模，相对于 Horovod (TensorFlow) 和 PyTorch 分布式数据并行，SMDDP 能够提供 20-40% 的吞吐量改进。 对于较小的集群大小和支持的模型，我们建议 SageMaker 培训编译器，这能够将整体工作时间减少多达 50%。

客户焦点：SageMaker 使用 Amazon Search 优化后端的 PyTorch Lightning

亚马逊搜索负责搜索和发现体验 Amazon.com. 它为在亚马逊上寻找产品的客户提供搜索体验。 在较高级别上，亚马逊搜索为在 Amazon.com 上销售的所有产品建立索引。 当客户输入查询时，亚马逊搜索会使用各种 ML 技术（包括深度学习模型）将相关且有趣的产品与客户查询进行匹配。 然后在向客户展示结果之前对产品进行排名。

亚马逊搜索科学家已使用 PyTorch Lightning 作为主要框架之一来训练支持搜索排名的深度学习模型，因为它在 PyTorch 之上增加了可用性功能。 在新的 SageMaker 发布之前，用 PyTorch Lightning 编写的深度学习模型不支持 SMDDP。 这阻止了喜欢使用 PyTorch Lightning 的 Amazon Search 科学家使用数据并行技术扩展他们的模型训练，显着减慢了他们的训练时间，并阻止他们测试需要更多可扩展训练的新实验。

该团队的早期基准测试结果显示，与单节点训练基线相比，在八个节点上训练样本模型的训练时间快 7.3 倍。 这些基准测试中使用的基线模型是一个多层感知器神经网络，具有七个密集的全连接层和超过 200 个参数。 下表总结了 ml.p3.16xlarge SageMaker 训练实例的基准测试结果。

接下来，我们深入了解新发布的细节。 如果你喜欢，你可以通过我们相应的 示例笔记本。

使用 SageMaker 分布式训练库运行 PyTorch Lightning

我们很高兴地宣布，SageMaker Data Parallel 现在与 SageMaker 培训中的 PyTorch Lightning 无缝集成。

PyTorch Lightning 是一个开源框架，可简化在 PyTorch 中编写自定义模型。 在某些方面类似于 Keras 为 TensorFlow 所做的，甚至可以说是 Hugging Face，PyTorch Lightning 为 PyTorch 本身的许多低级功能提供了一个带有抽象的高级 API。 这包括定义模型、分析、评估、修剪、模型并行、超参数配置、迁移学习等。

以前，PyTorch Lightning 开发人员不确定如何将他们的训练代码无缝迁移到高性能 SageMaker GPU 集群上。 此外，他们无法利用 SageMaker Data Parallel 带来的效率提升。

对于 PyTorch Lightning，一般来说，只需在 SageMaker Training 上运行这些 API，几乎不需要更改代码。 在示例笔记本中，我们使用 DDP策略 和 DDPP插件 方法。

使用 PyTorch Lightning 和 SageMaker Data Parallel 作为优化后端需要三个步骤：

1. 使用支持的 AWS 深度学习容器 (DLC) 作为您的基础镜像，或者选择创建您自己的容器和 自行安装 SageMaker Data Parallel 后端. 确保在必要的包中包含 PyTorch Lightning，例如 requirements.txt 文件中。
2. 对启用优化后端的训练脚本进行一些小的代码更改。 这些包括：
   1. 导入 SM DDP 库：

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. 为 SageMaker 设置 PyTorch Lightning 环境：

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. 如果您使用的 PyTorch Lightning 版本早于 1.5.10，则需要添加更多步骤。
      1. 首先，添加环境变量：

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. 二、确保你使用 DDPPlugin， 而不是 DDPStrategy. 如果您使用的是更新的版本，您可以通过放置 requirements.txt ，在 source_dir 对于您的工作，那么这不是必需的。 请参阅以下代码：

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. （可选）将您的进程组后端定义为 "smddp" ，在 DDPSTrategy 目的。 但是，如果您将 PyTorch Lightning 与 PyTorch DDP 后端一起使用，这也是受支持的，只需删除此 `process_group_backend`参数。 请参阅以下代码：

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. 确保您在估算器中注明了分布方法，例如 distribution={"smdistributed":{"dataparallel":{"enabled":True} 如果您使用的是 Herring 后端，或者 distribution={"pytorchddp":{"enabled":True}.

• 有关合适参数的完整列表 distribution 参数，请参阅我们的文档 点击此处.

现在您可以启动您的 SageMaker 培训工作了！ 您可以通过 Python SDK、Boto3、SageMaker 控制台、 AWS命令行界面 (AWS CLI) 以及无数其他方法。 从 AWS 的角度来看，这是 单个 API 命令: create-training-job. 无论您是从本地终端启动此命令， AWS Lambda 功能 亚马逊SageMaker Studio notebook、KubeFlow 管道或任何其他计算环境完全取决于您。

请注意，目前仅从 1.11 开始的较新版本的 PyTorch 支持 PyTorch Lightning 和 SageMaker Data Parallel 之间的集成。 此外，此版本仅适用于从 PyTorch 1.12 开始的 SageMaker 的 AWS DLC。 确保将此图像作为您的基础。 在 us-east-1，这个地址如下：

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

然后你可以使用它作为你的基础镜像来扩展你的 Docker 容器，或者你可以将它作为一个变量传递给 image_uri 的论点 SageMaker 训练估计器。

因此，您将能够在 SageMaker Training 的优化 GPU 上运行您的 PyTorch Lightning 代码，并在 AWS 上获得最佳性能。

在 SageMaker 上并行运行 PyTorch 分布式数据

最大的问题 PyTorch 分布式数据并行 (DDP) 解决看似简单：速度。 一个好的分布式训练框架应该提供稳定性、可靠性，最重要的是，在规模上提供出色的性能。 PyTorch DDP 通过为 Torch 开发人员提供 API 以在单节点和多节点设置中在多个 GPU 设备上复制他们的模型来实现这一点。 然后，该框架管理从训练数据集到每个模型副本的不同对象的分片，平均每个模型副本的梯度以在每个步骤中同步它们。 这会在完整的训练运行完全完成时生成一个模型。 下图说明了这个过程。

PyTorch DDP 在使用大型数据集的项目中很常见。 每个数据集的精确大小差异很大，但一般准则是以相似的比率缩放数据集、计算大小和模型大小。 也被称为 缩放定律，这三者的最佳组合非常有争议，并且会因应用而异。 在 AWS，基于与多个客户的合作，当整体数据集大小至少为几十 GB 时，我们可以清楚地看到数据并行策略的好处。 当数据集变得更大时，实施某种类型的数据并行策略是加快整个实验并缩短实现价值的关键技术。

以前，使用 PyTorch DDP 在本地或其他计算环境中进行分布式训练的客户缺乏一个框架来轻松地将他们的项目迁移到 SageMaker Training 上，以利用具有无缝控制平面的高性能 GPU。 具体来说，他们需要将他们的数据并行框架迁移到 SMDDP，或者在 SageMaker Training 上手动开发和测试 PyTorch DDP 的功能。 今天，SageMaker Training 很高兴为客户提供使用 PyTorch DDP 代码的无缝体验。

要有效地使用它，您无需对训练脚本进行任何更改。

您可以在以下代码中看到这个新参数。 在里面 distribution 参数，只需添加 pytorchddp 并将启用设置为 true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

此新配置从 SageMaker Python SDK 版本 2.102.0 和 PyTorch DLC 的 1.11 开始。

对于熟悉流行的 PyTorch DDP 开发人员 火炬运行 框架，知道这在 SageMaker 训练环境中不是必需的，它已经提供了强大的容错能力。 但是，为了最大程度地减少代码重写，您可以将另一个运行此命令的启动器脚本作为您的入口点。

现在 PyTorch 开发人员可以轻松地将他们的脚本移动到 SageMaker 上，确保他们的脚本和容器可以在多个计算环境中无缝运行。

这让他们准备在未来利用 SageMaker 的分布式训练库 提供 AWS 优化的训练拓扑以提供高达 40% 的加速增强。 对于 PyTorch 开发人员来说，这只是一行代码！ 对于 PyTorch DDP 代码，您可以简单地将后端设置为 smddp 在初始化（见 修改 PyTorch 训练脚本)，如以下代码所示：

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

正如我们在上面看到的，您还可以设置后端 DDPStrategy 至 smddp 使用闪电时。 这可能导致高达 整体加速 40% 对于大型集群！ 要了解有关 SageMaker 分布式培训的更多信息，请参阅 我们的点播网络研讨会，支持 笔记本电脑， 相关的 文件及 文件.

结论

在这篇文章中，我们介绍了 SageMaker 培训系列中的两个新功能。 这些使 PyTorch 开发人员更容易在 SageMaker 上使用他们现有的代码，包括 PyTorch DDP 和 PyTorch Lightning。

我们还展示了 Amazon Search 如何使用 SageMaker Training 来训练他们的深度学习模型，特别是 PyTorch Lightning 以及 SageMaker Data Parallel 优化的集体库作为后端。 总体而言，转向分布式培训帮助 Amazon Search 将培训时间缩短了 7.3 倍。

关于作者

艾米丽·韦伯（Emily Webber） 在 SageMaker 推出后就加入了 AWS，从那时起就一直在努力向全世界宣传它！ 除了为客户打造新的机器学习体验之外，Emily 还喜欢冥想和研究藏传佛教。

卡兰·迪曼 是 AWS 的一名软件开发工程师，常驻加拿大多伦多。 他对机器学习空间和构建用于加速分布式计算工作负载的解决方案充满热情。

维什瓦·卡里亚 是 AWS Deep Engine 的软件开发工程师。 她的兴趣在于机器学习和分布式系统的交叉点，她还热衷于在科技和人工智能领域赋予女性权力。

艾曼·埃尔纳拉维 是 Amazon Search 的首席软件工程师，负责机器学习加速、扩展和自动化方面的工作。 她的专长跨越多个领域，包括机器学习、分布式系统和个性化。