亚马逊-SageMaker-深度学习负载加速指南-全-

亚马逊 SageMaker 深度学习负载加速指南(全)

原文:annas-archive.org/md5/621c8e4bb6ef8d84ed646f8ded57a6f8

译者:飞龙

协议:CC BY-NC-SA 4.0

前言

深度学习(DL)是一种相对较新的机器学习类型,在自然语言理解和计算机视觉等任务中展现出了惊人的效果。有时,DL 的准确度甚至超过人类。

得益于开源框架的普及、公开可用的模型架构和预训练模型,许多人和组织能够成功地将前沿的深度学习模型应用于他们的实际用例。然而,开发、训练和部署深度学习模型还需要高度专业化且昂贵的硬件、软件栈、专业知识和管理能力,这可能会显著减缓其采用进程。

本书重点介绍如何在 Amazon SageMaker 上设计和管理深度学习工作负载,从而帮助您克服上述障碍。SageMaker 是一个广泛的 AWS 云机器学习平台,具有多种功能。本书并不打算详细介绍所有 SageMaker 功能,而是深入探讨与深度学习工作负载相关的特性。在编写本书时,我们优先考虑了深度而非广度。本书的目标是为您提供有关如何高效地在 Amazon SageMaker 上实施实时深度学习模型应用的实用指南。

由于云计算和机器学习的采用正在加速,本书可能对广泛的受众群体感兴趣,从初学者到经验丰富的机器学习从业者。具体来说,本书适合从事深度学习模型开发和训练的机器学习工程师,以及负责设计和优化深度学习工作负载的解决方案架构师。

假设您已经熟悉 Python 生态系统,以及机器学习和深度学习的基本原理。熟悉 AWS 并拥有与之合作的实践经验也是有帮助的。

随着章节的推进,从入门和概述主题到高级实现和优化技术,章节的复杂性逐步增加。您可以跳过某些章节,或者选择与您当前任务相关的具体主题。

本书的大多数章节都有相应的代码示例,您可以通过这些示例获得使用 Amazon SageMaker 的实际经验。建议您亲自尝试运行代码示例,当然,您也可以查看它们。我们还为每个章节的代码示例提供了注释。

请注意,运行代码示例会产生 AWS 费用。请确保查看 Amazon SageMaker 定价页面以获取详细信息。

我们欢迎您对本书的反馈和建议,并希望您享受您的学习旅程。

本书适用对象

本书是为已经具备深度学习(DL)领域基础知识,并希望学习和获得在 AWS 云上使用 Amazon SageMaker 服务功能训练和托管深度学习模型的实际经验的 DL 和 AI 工程师编写的。

本书涵盖的内容

第一章通过 Amazon SageMaker 介绍深度学习,将介绍 Amazon SageMaker:它如何简化基础设施和工作负载管理,以及该 AWS 服务的关键原则和主要功能。随后,我们将重点介绍托管训练、托管基础设施以及与其他 AWS 服务的集成。

第二章SageMaker 上的深度学习框架和容器,将详细回顾 SageMaker 如何广泛利用 Docker 容器。我们将首先深入了解流行深度学习框架(TensorFlow、PyTorch 和 MXNet)的预构建容器。接着,我们将探讨如何扩展预构建的 SageMaker 容器和使用自带容器。对于后者,我们将回顾在 SageMaker 中训练和服务容器的技术要求。

第三章SageMaker 开发环境管理,将讨论如何使用 CLI、SDK 和 CloudFormation 以编程方式管理 SageMaker 资源。我们将探讨如何使用 SageMaker Studio 和 Notebooks 组织高效的开发流程,以及如何与您喜欢的 IDE 集成。我们还将回顾如何使用 SageMaker Local Mode 排查深度学习代码的问题。我们将回顾 SageMaker 的各种功能,帮助我们组织和管理数据集,并讨论 AWS 上的各种存储选项及其应用场景。

第四章深度学习数据集管理,将提供关于在 SageMaker 上设置第一个深度学习项目的实用指南,然后构建、训练并使用一个简单的深度学习模型。我们将提供该项目的跟随实现,以便读者可以自行学习和实验 SageMaker 的核心功能。

第五章深度学习训练硬件考虑,将讨论最适合深度学习模型的实例的性价比特点,并覆盖在不同场景下如何选择不同实例类型以获得最佳性能。

第六章工程化分布式训练,将重点介绍分布式训练过程的常见方法,并解释为什么深度学习模型可能需要这样做。然后,我们将概述开源训练分发框架以及 SageMaker 原生分布式训练库。

第七章深度学习训练的运营化,将讨论如何使用 SageMaker Debugger 及其 Profiler 监控和调试深度学习训练任务,以及如何通过管理型 Spot 训练、提前停止和其他策略来优化成本。

第八章考虑硬件推理,将提供使用 PyTorch 和 Hugging Face 框架构建 NLP 最先进模型的实际指导。读者将跟随代码学习如何准备用于在 Amazon SageMaker 上进行分布式训练的训练脚本,并监控和进一步优化训练任务。我们将使用 SageMaker 数据并行库来分配训练计算。

第九章实现模型服务器,将从回顾 SageMaker 托管服务的关键组件开始,例如实时端点和批量推理任务、模型注册表和服务容器。读者将学习如何使用 Python SDK 配置其端点部署和批量推理任务。

第十章操作推理工作负载,将专注于 DL 服务器的软件堆栈,特别是模型服务器。我们将回顾 TensorFlow 和 PyTorch 解决方案提供的模型服务器,以及像 SageMaker 多模型服务器这样的框架无关模型服务器。我们将讨论何时选择某一选项而非另一选项。

为了充分利用本书

书中涵盖的软件/硬件 操作系统要求
已建立本地 SageMaker 兼容环境 Windows、macOS 或 Linux

如果您使用的是本书的数字版本,我们建议您自己输入代码或从本书的 GitHub 仓库获取代码(下一节提供了链接)。这样做将帮助您避免与复制和粘贴代码相关的潜在错误。

下载示例代码文件

您可以从 GitHub 下载本书的示例代码文件,网址为github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker。如果代码有更新,它将在 GitHub 仓库中更新。

我们还提供了来自我们丰富书籍和视频目录的其他代码包,网址为github.com/PacktPublishing/。快来看看吧!

下载彩色图片

我们还提供了一个 PDF 文件,里面包含了本书中使用的截图/图表的彩色图片。您可以在此下载:packt.link/FXLPc

使用的约定

本书中使用了多种文本约定。

代码文本:表示文本中的代码单词、数据库表名、文件夹名、文件名、文件扩展名、路径名、虚拟 URL、用户输入和 Twitter 用户名。以下是一个示例:“在以下代码块中,我们使用'_build_tf_config()'方法来设置此变量。”

一段代码块设置如下:

estimator.fit({
"train":"s3://unique/path/train_files/",
"test":"s3://unique/path/test_files"}
)

所有命令行输入或输出如下所示:

conda create -n sagemaker python=3.9

当我们希望引起你对代码块中特定部分的注意时,相关的行或项目会以粗体显示:

 tensorboard --logdir ${tb_debug_path}

粗体:表示新术语、重要词汇或屏幕上出现的单词。例如,菜单或对话框中的单词通常会以粗体显示。举个例子:“在创建默认的 IAM 角色弹出窗口中,选择任何 S3 存储桶。”

提示或重要说明

看起来像这样。

与我们联系

我们始终欢迎读者的反馈。

常规反馈:如果你对本书的任何内容有疑问,请在邮件主题中注明书名,并通过 customercare@packtpub.com 联系我们。

勘误表:虽然我们已经尽力确保内容的准确性,但仍然可能会有错误。如果你发现本书中有错误,我们将非常感谢你向我们报告。请访问 www.packtpub.com/support/errata,选择你的书籍,点击勘误提交表单链接并填写相关细节。

盗版:如果你在网上发现我们作品的任何非法复制品,无论何种形式,我们将非常感谢你提供该材料的位置地址或网站名称。请通过 copyright@packt.com 联系我们并附上该材料的链接。

如果你有兴趣成为作者:如果你在某个领域拥有专业知识,并且有兴趣撰写或贡献一本书,请访问 authors.packtpub.com

分享你的想法

阅读完 《使用 Amazon SageMaker 加速深度学习工作负载》 后,我们很想听听你的想法!请 点击这里直接访问亚马逊的书籍评价页面,并分享你的反馈。

你的评价对我们以及技术社区至关重要,将帮助我们确保提供高质量的内容。

下载本书的免费 PDF 副本

感谢购买本书!

你喜欢随时随地阅读,但又无法随身携带纸质书籍吗?你的电子书购买无法兼容你选择的设备吗?

别担心,现在购买每本 Packt 图书,你都可以免费获得该书的无 DRM PDF 版本。

在任何地方、任何设备上阅读。可以直接搜索、复制并粘贴代码到你最喜欢的技术书籍中。

优惠不仅仅到此为止,你还可以获得独家折扣、新闻通讯以及每天直接发送到邮箱的精彩免费内容。

按照这些简单的步骤来获取优惠:

  1. 扫描二维码或访问下面的链接

https://packt.link/free-ebook/9781801816441

  1. 提交你的购买证明

  2. 就这些!我们会直接将免费的 PDF 和其他福利发送到你的电子邮件中

第一部分:在 Amazon SageMaker 上的深度学习简介

在第一部分,我们将从简要介绍深度学习和 Amazon SageMaker 开始,接着重点讲解在本书中将使用的 SageMaker 关键功能。

本节包含以下章节:

  • 第一章通过 Amazon SageMaker 引入深度学习

  • 第二章SageMaker 上的深度学习框架和容器

  • 第三章管理 SageMaker 开发环境

  • 第四章管理深度学习数据集

第一章:介绍与 Amazon SageMaker 的深度学习

深度学习DL)是一个相对较新的但正在积极发展的机器学习ML)领域。在过去的 15 年里,深度学习从研究实验室走向了我们的家庭(如智能家居和智能音响)、汽车(即自动驾驶功能)、手机(例如,照片增强软件)以及你每天使用的应用程序(如你最喜欢的视频平台中的推荐系统)。

深度学习模型在计算机视觉(如物体检测与分割、图像分类任务和图像生成)以及语言任务(如翻译、实体提取和文本情感分析)等任务中,已经达到了甚至超越了人类的准确性。除这些领域外,深度学习还被广泛应用于复杂领域,如医疗保健、信息安全、机器人技术和自动化。

我们应该预期这些领域中的深度学习应用将随着时间的推移不断增长。随着当前成果和未来的承诺,也伴随而来的是实现深度学习模型时的挑战。但在讨论这些挑战之前,我们先来快速回顾一下深度学习是什么。

在本章中,我们将完成以下内容:

  • 我们将快速回顾一下深度学习及其挑战

  • 我们将概述 Amazon SageMaker 以及它在深度学习项目中的价值主张

  • 我们将概述 SageMaker 的基础组件——即托管训练和托管栈

  • 我们将概述其他关键的 AWS 服务

以下主题将涵盖:

  • 使用 Amazon SageMaker 探索深度学习

  • 选择 Amazon SageMaker 处理深度学习工作负载

  • 探索 SageMaker 的托管训练栈

  • 使用 SageMaker 的托管托管栈

  • 与 AWS 服务的集成

技术要求

本章中有多个动手编码示例。要跟随这些示例,您需要以下内容:

本章中的所有代码可以从 github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker 下载。

使用 Amazon SageMaker 探索深度学习

深度学习是机器学习领域的一个子集,它使用一种特定的架构:各层之间相互连接的可学习参数。在这种架构中,每一层都在从训练数据集中“学习”一种表示。每一组新的训练数据样本会稍微调整模型各层中的可学习参数,以最小化损失函数。堆叠的层数构成了模型的“深度”。在推理时(即,当我们使用模型从输入信号推理输出时),每一层接收来自前一层输出的输入,根据输入计算其表示,并将其传递给下一层。

图 1.1 – 全连接深度学习网络

图 1.1 – 全连接深度学习网络

简单的深度学习模型可能仅由几个全连接层组成,而最先进的SOTA)模型则有数百个层次,包含数百万甚至数十亿个可学习的参数。而且,模型的规模还在不断增长。例如,我们来看看 GPT 系列模型在各种自然语言处理任务中的演变。GPT-1 模型于 2018 年发布,拥有 1.1 亿个参数;2019 年发布的 GPT-2 模型拥有 15 亿个参数;最新版本 GPT-3 于 2020 年发布,拥有 1750 亿个参数!

随着参数数量的增加,深度学习从业者面临着若干工程问题:

  • 我们如何在训练时将模型放入实例内存?如果不可能,那么如何将模型拆分到多个 GPU 设备和/或计算节点的内存中?

  • 我们如何在训练时组织多个节点之间的通信,以便整体模型能够汇总各个节点的学习成果?

层内部结构也变得越来越复杂,需要更多的计算能力。深度学习模型通常还需要大量特定格式的数据。

因此,为了能够成功地训练和使用最先进的深度学习模型,机器学习工程师需要解决以下任务(当然,除了实现最先进的模型):

  • 在训练和推理过程中,获得大量专用计算资源

  • 设置并维护软件堆栈(例如,GPU 库、深度学习框架和加速库)

  • 实现、管理和优化分布式训练任务

  • 实现、部署和监控推理管道

  • 在调优模型性能时,组织高效的时间和成本实验

  • 预处理、标记和访问大数据集(数 GB 和 TB 的数据)

如你所见,这些任务不一定与解决特定的业务问题相关。然而,你需要将这些组件都配置正确,以确保能够使用深度学习模型在时间和预算的范围内,以最高的准确度解决某一特定的业务问题。

使用 SageMaker

Amazon SageMaker 是 AWS 的一项服务,旨在通过消除“无差异的繁重工作”(例如前面提到的任务)来简化机器学习从业者的工作,让你能专注于实际解决业务问题。它集成了各种功能,帮助构建、训练和部署机器学习模型。SageMaker 首次推出是在 2017 年末,之后大幅扩展,仅在 2020 年,根据 AWS 的Re:invent 2020大会的主题演讲,便新增了 200 多项功能。

Amazon SageMaker 面向广泛的受众和使用场景,包括以下内容(这只是一个非独占的列表):

  • 没有太多机器学习背景的开发者

  • 企业数据科学团队

  • 领先的研究机构希望开展前沿研究

SageMaker 是一个托管服务,通过 API 抽象了底层计算资源和软件组件的管理。AWS 客户使用这些 API 创建训练任务,管理模型工件,并将模型部署到 Amazon SageMaker 进行推理。AWS 负责 SageMaker 资源的高可用性,并提供相应的服务级别协议SLAs)。Amazon SageMaker 采用“按需付费”模式,客户仅为实际使用的资源付费。

在本书中,我们将探索与深度学习模型和工作负载相关的 SageMaker 功能,并为流行的深度学习用例(如自然语言处理NLP)和计算机视觉CV))构建端到端解决方案。

我们将专注于以下 Amazon SageMaker 功能:

  • 模型开发阶段:

    • 使用 SageMaker GroundTruth 进行数据准备

    • 使用 SageMaker Processing 进行数据的预处理和后处理

    • 与数据存储解决方案的集成 – Amazon S3、Amazon EFS 和 Amazon FSx for Lustre

    • 使用 SageMaker Studio IDE 和 Notebooks 开发模型

  • 训练阶段:

    • 托管的训练实例和软件堆栈

    • 支持 TensorFlow 和 Pytorch 框架的深度学习容器

    • 使用 SageMaker 的 DataParallel 和 ModelParallel 库实现分布式训练

    • 使用 SageMaker Debugger 进行训练监控和调试

  • 推理:

    • 托管的批处理和实时推理平台

    • 在推理时对模型进行监控

    • 为深度学习服务计算实例

在本书的进程中,我们还将学习如何使用 SageMaker API 和 SDK 来编程管理我们的资源。此外,我们还将讨论优化策略和 SageMaker 功能,以降低成本并缩短上市时间。

本书专注于深度学习功能,因此我们将暂时搁置一些 SageMaker 的功能和特点。你可以通过阅读 SageMaker 文档(docs.aws.amazon.com/sagemaker/latest/dg/whatis.xhtml)和实践博客(aws.amazon.com/blogs/machine-learning/category/artificial-intelligence/sagemaker/)进一步探索 Amazon SageMaker。

在接下来的章节中,我们将了解 SageMaker 如何帮助我们处理深度学习工作负载。

选择 Amazon SageMaker 进行深度学习工作负载

如前所述,深度学习工作负载由于需要访问大量专门资源(主要是 GPU 设备和高吞吐量存储解决方案),因此面临多种工程挑战。然而,管理软件栈也是一个挑战,因为机器学习和深度学习框架会频繁发布新版本。由于相关成本较高,因此必须高效地组织训练和推理工作,以避免浪费。

让我们来回顾一下 SageMaker 如何应对这些挑战。

托管的计算和存储基础设施

SageMaker 为你的训练和推理工作负载提供了完全托管的计算基础设施。SageMaker 训练和推理集群能够在几分钟内扩展至几十个甚至上百个独立实例。这在需要快速访问大型计算集群且只需要短时间使用的场景中尤其有用(例如,你需要每隔几个月就训练一次复杂的深度学习模型,并且数据集很大)。与其他 AWS 服务一样,SageMaker 资源为推理端点提供了先进的自动扩展功能,以便客户根据需求调整资源,而无需过度配置资源。

你还可以根据特定深度学习模型和工作负载类型的需求,从越来越多的计算实例中进行选择。例如,在许多场景下,你可能需要使用基于 GPU 的实例来训练你的深度学习模型,而在推理时,可能可以使用较便宜的 CPU 实例,而不会影响最终用户性能。SageMaker 使你能够根据特定的深度学习模型和任务选择最优实例。

在发生故障时,AWS 会自动替换故障实例,无需客户干预。

这些 SageMaker 功能极大地惠及客户,因为 SageMaker 简化了机器学习基础设施的容量规划和运营管理。

管理的深度学习软件栈

要构建、训练和部署深度学习模型,你需要使用各种框架和软件组件来执行专门的计算以及在分布式集群中进行设备和节点之间的通信。在不同开发环境中创建和维护软件栈可能会非常繁琐。

为了解决这些需求,作为 SageMaker 生态系统的一部分,AWS 提供了多个预构建的开源 Docker 容器,支持流行的深度学习框架,如 PyTorch、TensorFlow、MXNet 等。这些容器由 AWS 构建和测试,并针对特定任务(例如,训练与推理的不同容器)和计算平台(基于 CPU 或 GPU 的容器,不同版本的 CUDA 工具包等)进行了优化。

由于 Docker 容器提供了互操作性和封装性,开发人员可以利用预构建的 SageMaker 容器,在将工作负载部署到云集群之前,先在本地构建和调试它们,从而缩短开发周期。你还可以根据特定需求扩展或修改 SageMaker 容器。

高级操作功能

虽然 Amazon SageMaker 利用多个流行的开源深度学习解决方案,但它还提供了一些独特的功能来解决在实现 ML 工作负载时遇到的特定挑战,例如:

  • SageMaker Debugger

  • SageMaker 模型监控

  • SageMaker 的 DataParallel/ModelParallel 分布式训练库

接下来让我们继续查看与其他 AWS 服务的集成。

与其他 AWS 服务的集成

Amazon SageMaker 与其他 AWS 服务紧密集成,帮助开发人员构建可扩展、高效且安全的工作负载。

在下一部分,我们将看看如何利用 Amazon SageMaker 的托管训练堆栈来运行深度学习模型。

探索 SageMaker 的托管训练堆栈

Amazon SageMaker 提供一套功能和与其他 AWS 服务的集成点,用于配置、运行和监控机器学习训练任务。通过 SageMaker 托管训练,开发人员可以执行以下操作:

  • 从各种内置算法和容器中选择,或者使用自带的模型(BYO 模型)

  • 根据模型需求,从广泛的计算实例中选择

  • 使用 SageMaker Debugger 调试和分析其训练过程,接近实时

  • 运行偏差检测和模型可解释性任务

  • 运行增量训练任务,从检查点恢复并使用 Spot 实例

Spot 实例

Amazon EC2 Spot 实例为客户提供了以较低价格(最高可达 90%)访问未使用的计算资源的机会。当其他人抢占这些实例时,Spot 实例会被释放,导致工作负载中断。

  • 运行模型调优任务,搜索最优的模型超参数组合

  • 在可搜索的实验目录中组织训练任务

Amazon SageMaker 提供以下开箱即用的功能:

  • 配置、引导和销毁训练节点

  • 捕获训练任务的日志和指标

在本节中,我们将逐步介绍 SageMaker 训练任务的所有阶段及其相关组件。请参考下图,图中提供了创建、管理和监控您第一个 SageMaker 训练任务的逐步指南。我们将在本书的 第二部分 讨论 SageMaker 管理训练堆栈的高级功能:

图 1.2 – Amazon SageMaker 训练架构

图 1.2 – Amazon SageMaker 训练架构

让我们逐步了解每个步骤。

我们还将提供代码示例,说明如何使用 SageMaker Python SDK 配置 SageMaker 训练任务(sagemaker.readthedocs.io/en/stable/)。

步骤 1 – 配置并创建训练任务

您可以通过 API 调用实例化 SageMaker 训练任务。

SageMaker 定义了几个必须提供的配置参数。它们如下所示。

选择训练算法

Amazon SageMaker 支持多种类型的机器学习算法:

  • 内置算法是所有 SageMaker 用户开箱即用的。截止到本文写作时,18 个内置算法涵盖了各种使用场景,包括计算机视觉和自然语言处理任务的深度学习算法。用户只需要负责提供算法的超参数。

  • 自定义算法由用户开发。在这种情况下,AWS 不负责训练逻辑。训练脚本将在 Docker 容器内执行。开发人员可以选择使用 AWS 提供的预装软件依赖的 Docker 镜像,也可以使用自带 Docker 镜像。

  • 市场算法由第三方供应商开发,并通过 AWS Marketplace 提供。与内置算法类似,它们通常提供完全托管的体验,用户负责提供算法的超参数。不同于内置算法是免费的,使用市场算法时用户通常需要支付费用。

定义 IAM 角色

Amazon SageMaker 依赖于 Amazon IAM 服务,特别是 IAM 角色,用于定义哪些 AWS 资源和服务可以从训练任务中访问。这就是为什么在每次调度 SageMaker 训练任务时,都需要提供一个 IAM 角色,之后该角色会被分配给训练节点。

定义训练集群

另一组必需的参数定义了训练集群的硬件配置,包括多个计算实例、实例类型以及实例存储。

建议根据具体需求仔细选择实例类型。至少,ML 工程师需要了解训练时使用的是哪种计算设备。例如,在大多数深度学习(DL)模型的情况下,通常需要使用基于 GPU 的实例,而许多经典的机器学习算法(如线性回归或随机森林)则是 CPU 绑定型的。

机器学习工程师还需要考虑应该配置多少实例。在配置多个节点时,您需要确保您的算法和训练脚本支持分布式训练。

内置算法通常会在其公共文档中提供推荐的实例类型和数量。它们还会定义是否支持分布式训练。如果不支持分布式训练,您应该配置单节点训练集群。

定义训练数据

Amazon SageMaker 支持几种用于训练数据的存储解决方案:

  • Amazon S3:这是一个低成本、高耐久性且高可用性的对象存储。它被视为存储训练数据集的默认选择。Amazon S3 支持两种输入模式(也在训练作业配置中定义)用于训练数据集:

    • 文件:Amazon SageMaker 将训练数据集从 S3 位置复制到本地目录

    • 管道:Amazon SageMaker 通过 Unix 命名管道直接从 S3 流式传输数据到容器

    • FastFile:Amazon SageMaker 提供的一项新的文件流能力。

  • Amazon EFS:这是一个弹性文件系统服务。如果用于持久化训练数据,Amazon SageMaker 会自动将训练实例挂载到共享文件系统。

  • Amazon FSx for Lustre:这是一个高性能的共享文件系统,优化了最低延迟和最高吞吐量。

在训练开始之前,您需要确保数据已经存储在这些解决方案中的某个位置,然后提供数据集的位置。

请注意,您可以在训练作业中提供多个数据集的位置(例如,训练集、测试集和评估集)。

选择您的算法超参数

虽然这不是严格要求的,但在大多数情况下,您需要定义算法的某些超参数。此类超参数的示例包括批量大小、训练轮数和学习率。

在训练时,这些超参数将作为命令行参数传递给训练脚本。在自定义算法的情况下,开发人员需要负责在训练脚本中解析并设置超参数。

定义训练度量指标

度量指标是另一个可选但重要的参数。SageMaker 提供了与 Amazon CloudWatch 的开箱即用集成,可以实时传输训练日志和度量指标。在日志的情况下,SageMaker 会自动将 stdoutstderr 从训练容器传输到 CloudWatch。

stdout 和 stderr

stdoutstderr 是 Linux 和类 Unix 操作系统中的标准数据流。每次运行 Linux 命令时,这些数据流都会自动建立。正常的命令输出会发送到 stdout;任何错误消息则会发送到 stderr

对于指标,用户需要首先定义每个指标的正则表达式模式。在训练时,运行在训练实例上的 SageMaker 实用程序将监控stdoutstderr中的正则表达式模式匹配,然后提取该指标的值并将指标名称和值提交到 CloudWatch。因此,开发人员可以在 CloudWatch 中近实时地监控训练过程。

一些常见的训练指标示例包括损失值和准确性度量。

配置 SageMaker 训练任务以进行图像分类

在下面的 Python 代码示例中,我们将演示如何配置一个简单的训练任务,使用内置的图像分类算法(docs.aws.amazon.com/sagemaker/latest/dg/image-classification.xhtml):

  1. 从您的初始导入开始:

    import sagemaker
    from sagemaker import get_execution_role
    
  2. get_execution_role()方法允许您获取当前的 IAM 角色。此角色将用于调用 SageMaker API,而sagemaker.Session()则存储与 SageMaker 及其他 AWS 服务(如 S3)的交互上下文:

    role = get_execution_role() 
    sess = sagemaker.Session() 
    
  3. .image_uris.retrieve()方法允许您识别与内置图像分类算法相关的正确容器。请注意,如果选择使用自定义容器,您需要为您的特定训练容器指定一个 URI:

    training_image = sagemaker.image_uris.retrieve('image-classification', sess.boto_region_name)
    
  4. 定义训练集群中的实例数量。由于图像分类支持分布式训练,我们可以为训练集群分配多个实例以加速训练:

    num_instances = 2
    
  5. 图像分类算法需要基于 GPU 的实例,因此我们将选择使用 SageMaker P2 实例类型:

    instance_type = "ml.p2.xlarge"
    
  6. 接下来,我们必须定义训练和验证数据集的位置。请注意,图像分类算法支持多种数据格式。在本例中,我们选择使用 JPG 文件格式,这还需要.lst文件列出所有可用的图像:

    data_channels = {
        'train': f"s3://{sess.default_bucket()}/data/train", 
        'validation': f"s3://{sess.default_bucket()}/data/validation", 
        'train_lst': f"s3://{sess.default_bucket()}/data/train.lst",
        'vadidation_lst': f"s3://{sess.default_bucket()}/data/validation.lst",
    }
    
  7. 配置训练的超参数:

    hyperparameters=dict(
        use_pretrained_model=1,
        image_shape='3,224,224',
        num_classes=10,
        num_training_samples=40000, # TODO: update it
        learning_rate=0.001,
        mini_batch_size= 8    
    )
    
  8. 配置Estimator对象,它封装了训练任务的配置:

    image_classifier = sagemaker.estimator.Estimator(
        training_image,
        role, 
        train_instance_count= num_instances, 
        train_instance_type= instance_type,
        sagemaker_session=sess,
        hyperparameters=hyperparameters,
    )
    
  9. fit()方法将训练任务提交给 SageMaker API。如果没有问题,您应该会在 AWS 控制台中看到一个新的训练任务实例。您可以通过进入Amazon SageMaker | Training | Training Jobs来查看:

    image_classifier.fit(inputs=data_channels, job_name="sample-train")
    

接下来,我们将配置训练集群。

第 2 步 – 配置 SageMaker 训练集群

提交训练任务请求后,SageMaker 会自动执行以下操作:

  • 分配请求的训练节点数量

  • 分配 Amazon EBS 卷并将其挂载到训练节点上

  • 为每个节点分配 IAM 角色

  • 启动各种实用工具(如 Docker、SageMaker 工具包库等)

  • 将训练配置(超参数、输入数据配置等)定义为环境变量

接下来是训练数据。

第 3 步 – SageMaker 访问训练数据

当您的训练集群准备就绪时,SageMaker 为计算实例建立访问训练数据的权限。访问训练数据的确切机制取决于您的存储解决方案:

  • 如果数据存储在 S3 中并且输入模式为File,则数据将下载到实例的 EBS 卷上。请注意,根据数据集的大小,下载数据可能需要几分钟。

  • 如果数据存储在 S3 中并且输入模式为Pipe,则数据将根据需要在训练时从 S3 流式传输。

  • 如果数据存储在 S3 中,并且输入模式为FastFile,则训练程序将像文件存储在训练节点上一样访问这些文件。但在底层,文件将从 S3 流式传输。

  • 如果数据存储在 EFS 或 FSx for Luster 中,则训练节点将挂载在文件系统上。

训练继续部署容器。

第 4 步 - SageMaker 部署训练容器

SageMaker 会自动从 ECR 仓库中拉取训练映像。请注意,内置算法会抽象出底层的训练映像,使用户无需定义容器映像,只需定义要使用的算法。

第 5 步 - SageMaker 启动并监控训练作业

要启动训练作业,SageMaker 在所有训练节点上执行以下命令:

docker run [TrainingImage] train

如果训练集群具有带 GPU 设备的实例,则将使用nvidia-docker

训练脚本启动后,SageMaker 会执行以下操作:

  • 捕获stdout/stderr并将其发送到 CloudWatch 日志。

  • 运行正则表达式模式匹配以获取度量指标并将度量值发送到 CloudWatch。

  • 从 SageMaker API 监听SIGTERM信号(例如,如果用户决定提前停止训练作业)。

  • 监控是否发生早停止条件并在此情况下发出SIGTERM

  • 监控训练脚本的退出代码。在非零退出代码的情况下,SageMaker 将标记训练作业为“失败”。

第 6 步 - SageMaker 将持久化训练产物

无论训练作业成功还是失败,SageMaker 都会将产物存储在以下位置:

  • /opt/ml/output目录,可以用来在作业完成后持久化任何训练产物。

  • /opt/ml/model目录,其内容将被压缩成.tar格式并存储在 SageMaker 模型注册表中。

一旦您的第一个模型训练完毕以解决特定的业务问题,下一步就是使用您的模型(在 ML 术语中,进行推断)。在接下来的几节中,我们将了解 SageMaker 为运行各种用例的 ML 推断工作负载提供了哪些能力。

使用 SageMaker 的托管主机堆栈

Amazon SageMaker 支持多种类型的托管主机基础设施:

  • 一个持久的同步 HTTPS 端点,用于实时推断

  • 一个用于准实时推断的异步端点

  • 一个瞬时的批量转换作业,用于整个数据集的推断

在下一部分,我们将讨论在何时使用哪种类型的托管基础设施的用例,并将详细回顾实时推理端点。

实时推理端点。

实时端点专为需要尽快获得推理结果的用例设计。SageMaker 的实时端点是一个 HTTPS 端点:模型输入由客户端通过 POST 请求负载提供,推理结果在响应体中返回。通信是同步的。

有许多场景适合使用实时端点,例如:

  • 基于用户的观看历史、个人评分和当前流行趋势,在用户打开流媒体应用程序时提供电影推荐。

  • 检测实时视频流中的物体。

  • 在用户输入文本时生成建议的下一个单词。

SageMaker 实时端点为客户提供一系列功能,用于设计和管理其推理工作负载:

  • 创建一个完全托管的计算基础设施,并实现横向扩展(意味着单个端点可以使用多个计算实例来应对高流量负载,而不会导致性能下降)。

  • 有多种 EC2 计算实例类型可供选择,具体取决于模型的需求,包括 AWS 的自定义芯片 Inferentia 和 SageMaker 弹性推理。

  • 针对流行的深度学习框架提供预构建的推理容器。

  • 多模型和多容器端点。

  • 用于 A/B 测试的模型生产变体。

  • 多模型推理管道。

  • 性能、准确性和偏差的模型监控。

  • 使用 SageMaker Neo 和 SageMaker Edge Manager 在边缘设备上优化和管理推理。

由于这是一个托管能力,Amazon SageMaker 负责管理用户实时端点的以下方面:

  • 根据客户定义的扩展策略提供和扩展基础计算基础设施。

  • 在单个 SageMaker 端点上部署多个模型时,模型版本和容器之间的流量塑形。

  • 在计算实例和模型级别流式传输日志和度量指标。

创建和使用您的 SageMaker 端点。

让我们一起走过配置、提供和使用第一个 SageMaker 实时端点的过程。这将帮助您理解其内部工作原理和可用的配置选项。以下图表提供了一个视觉指南:

图 1.3 – SageMaker 推理端点的部署与使用

图 1.3 – SageMaker 推理端点的部署与使用。

第一步 – 启动端点创建。

有几种方法可以启动 SageMaker 端点创建:SageMaker Python SDK、boto3 SDK、AWS CLI 或通过 CloudFormation 模板。在请求过程中,您需要提供以下几个参数:

  • SageMaker Model Registry 中的模型定义,将在推理时使用。模型定义包括对 S3 中序列化模型工件的引用,以及对 Amazon ECR 中推理容器(或多容器端点的情况下的多个容器)的引用。

  • 端点配置,定义计算实例的数量和类型,以及(可选)多个模型的组合(在多模型端点的情况下)或多个模型生产变体的组合(在 A/B 测试的情况下)。

第二步 – 配置 SageMaker 端点用于图像分类

以下 Python 代码示例展示了如何使用先前训练的图像分类模型创建和部署端点:

  1. 从初始导入、IAM 角色和 SageMaker 会话实例化开始:

    import sagemaker
    from sagemaker import get_execution_role
    role = get_execution_role()
    sess = sagemaker.Session()
    
  2. 获取图像分类算法的推理容器 URI:

    image_uri = sagemaker.image_uris.retrieve('image-classification', sess.boto_region_name)
    
  3. 定义存储在 S3 中的模型工件(如训练权重)的位置:

    model_data = f"s3://{sess.default_bucket}/model_location"
    
  4. 创建一个 SageMaker Model对象,封装模型配置:

    model = Model(
    image_uri=image_uri, 
    model_data=model_data,
    name="image-classification-endpoint",
    sagemaker_session=sess,
    role=role
    )
    
  5. 定义端点配置参数:

    endpoint_name = "image-classification-endpoint"
    instance_type = "ml.g4dn.xlarge"
    instance_count = 1
    
  6. .predict()方法向 SageMaker 提交请求,创建部署了特定模型的端点:

    predictor = model.deploy(
    instance_type=instance_type, 
    initial_instance_count=instance_count,
    endpoint_name=endpoint_name,
    )
    

完成这些步骤后,SageMaker 开始工作。

第三步 – SageMaker 配置端点

提交配置请求后,SageMaker 执行以下操作:

  • 根据端点配置分配多个实例

  • 它将部署推理容器

  • 它将下载模型工件

从开始到完成端点配置需要几分钟时间。配置时间取决于多个参数,如实例类型、推理容器大小以及需要上传到推理实例的模型工件大小。

请注意,SageMaker 不直接暴露推理实例。相反,它使用前端负载均衡器,然后在配置的实例之间分发流量。作为托管服务,您永远不会直接与推理实例交互,只能通过 SageMaker API 进行操作。

第四步 – SageMaker 启动模型服务器

一旦端点完全配置完成,SageMaker 通过运行以下命令启动推理容器,执行容器中的ENTRYPOINT命令:

docker run image serve

此脚本执行以下操作:

  • 启动模型服务器,暴露 HTTP 端点

  • 使模型服务器将模型工件加载到内存中

  • 在推理时,使模型服务器执行推理脚本,定义如何预处理数据

在 SageMaker 托管的 Docker 镜像情况下,模型服务器和启动逻辑已由 AWS 实现。如果选择自己带服务的容器(BYO serving container),则需要单独实现。

SageMaker 捕获 stdout/stderr 流并自动将其流式传输到 CloudWatch 日志。它还会流式传输实例指标,如总调用次数和每个实例的调用次数、调用错误和延迟指标。

步骤 5 – SageMaker 端点处理流量

一旦模型服务器启动并运行,最终用户可以向 SageMaker 端点发送 POST 请求。该端点会根据授权头进行请求授权(当使用 SageMaker Python SDK 或 AWS CLI 时,这些头部会根据 IAM 配置文件自动生成)。如果授权成功,则将负载发送到推理实例。

运行中的模型服务器通过执行推理脚本来处理请求,并返回响应负载,然后将其传递给最终用户。

步骤 6 – SageMaker 扩展推理端点的规模

您可以选择定义自动扩展策略,以便在端点实例之间进行扩展或缩减。在这种情况下,SageMaker 会根据需求更高效地添加或移除端点后面的计算节点。请注意,SageMaker 仅支持水平扩展,例如添加或移除计算节点,而不支持更改实例类型。

SageMaker 支持多种类型的扩展事件:

  • 手动方式,用户通过 API 调用更新端点配置

  • 一种目标追踪策略,SageMaker 根据用户定义的指标的值(例如调用次数或资源利用率)来扩展或缩减

  • 一种步骤扩展策略,它提供更细粒度的控制,允许用户根据阈值的突破程度调整实例的数量

  • 一种计划扩展策略,允许您根据特定的计划扩展 SageMaker 端点(例如,在流量较低的周末进行缩减,在流量较高的工作日进行扩展)

高级模型部署模式

我们刚刚回顾了一个简单的单模型实时端点的结构。然而,在许多现实场景中,需要在任何给定时间点都能提供数十个或数百个模型,这种方法将导致大量的推理节点未被充分利用或不均匀使用。

这种情况通常是不可取的,因为它会导致高昂的计算成本而没有为最终用户带来任何价值。为了解决这个问题,Amazon SageMaker 提供了一些高级部署选项,允许您将多个模型合并到同一个实时端点中,从而更高效地利用资源。

多容器端点

部署多容器端点时,您可以在同一个端点中指定最多 15 个不同的容器。每个推理容器都有自己的模型工件和运行时环境。这使得您能够在单个 SageMaker 端点内部署使用不同框架和运行时环境构建的模型。

在创建时,你需要定义一个唯一的容器主机名。然后,每个容器可以独立调用。在终端调用期间,你需要将该容器主机名作为请求头之一提供。SageMaker 将根据该请求头自动将推理请求路由到正确的容器。

当有多个模型且流量较低,且这些模型有不同的运行时环境(例如,Pytorch 和 TensorFlow)时,这个功能非常实用。

推理管道

类似于多容器终端,推理管道允许你在单个 SageMaker 终端中组合不同的模型和容器运行时环境。然而,这些容器是按顺序调用的。该功能适用于需要使用不同运行时要求进行前处理和/或后处理的推理请求场景;例如:

  • 前处理阶段使用 scikit-learn 库完成

  • 推理使用深度学习框架完成

  • 后处理使用自定义运行时环境完成,例如 Java 或 C++

通过将推理管道的不同阶段封装在单独的容器中,一个容器的变化不会对其他容器产生不利影响,例如更新依赖版本。由于推理管道中的容器位于同一计算节点上,因此可以保证容器之间请求交接时的低延迟。

多模型终端

多模型终端允许你在单个终端中部署数百个模型。与多容器终端和推理管道不同,多模型终端只有一个运行时环境。SageMaker 会自动将模型工件加载到内存中,并处理推理请求。当模型不再需要时,SageMaker 会将其从内存中卸载,以释放资源。这会导致在一段时间后首次调用模型时出现一些额外的延迟。模型工件存储在 Amazon S3 中,并由 SageMaker 自动加载。

多模型终端的核心是 AWS 开发的开源 Multi-Model Server,它提供模型管理功能(加载、卸载和资源分配),并通过 HTTP 前端接收推理请求,执行给定模型的推理代码,并返回结果负载。

多模型终端在有大量同质模型且最终用户可以容忍预热延迟时最为适用。

SageMaker 异步终端

到目前为止,我们讨论了 SageMaker 实时终端,它们是同步工作的:用户通过发送 POST 请求调用终端,等待终端运行推理代码,然后返回推理结果到响应负载中。推理代码预计会在 60 秒内完成;否则,SageMaker 终端将返回超时响应。

然而,在某些场景中,这种同步通信模式可能会带来问题:

  • 大型模型可能需要相当长的时间来进行推理

  • 大型负载数据(例如,高分辨率图像)

对于此类场景,SageMaker 提供了异步端点,允许您排队推理请求并异步处理,避免了潜在的超时问题。异步端点还允许处理更大的负载,最大可达 1 GB,而 SageMaker 实时端点的限制为 5 MB。异步端点可以在推理队列为空时缩减到 0 个实例,以提供额外的成本节省。这对于具有零星推理流量模式的场景特别有用。

异步端点的主要权衡是推理结果在接近实时的情况下交付,可能不适合那些对一致延迟有要求的场景:

图 1.4 – SageMaker 异步端点

图 1.4 – SageMaker 异步端点

SageMaker 批量转换

SageMaker 批量转换允许您对一批推理输入进行预测。这对于那些有重复业务流程且对延迟没有严格要求的场景非常有用。例如,每晚执行的任务用于计算负载应用程序的风险。

SageMaker 批量转换对以下用例非常有利:

  • 客户仅为任务执行期间消耗的资源付费

  • 批量转换任务可以扩展到 GB 级和数十个计算节点

在调度批量转换任务时,您需要定义集群配置(计算节点的类型和数量)、模型工件、推理容器、推理数据集的输入 S3 位置以及生成的预测结果的输出 S3 位置。请注意,客户可以使用相同的容器用于 SageMaker 实时端点和批量转换任务。这使得开发人员可以使用相同的模型/容器进行在线预测(作为实时端点)和离线预测(作为批量转换任务):

图 1.5 – SageMaker 批量转换任务

图 1.5 – SageMaker 批量转换任务

通过此方法,您可以了解如何使用 SageMaker 训练任务训练一个简单的深度学习模型,并创建一个实时端点进行推理。在继续之前,我们需要了解一些亚马逊 SageMaker 使用的基础 AWS 服务,这些服务将贯穿本书。

与 AWS 服务的集成

亚马逊 SageMaker 依赖于多个 AWS 服务,如存储和密钥管理。在本节中,我们将回顾与其他 AWS 服务的关键集成,以及它们在何种场景下可能有用。

数据存储服务

数据存储服务是构建任何机器学习工作负载的关键。AWS 提供了多种存储解决方案,以应对各种实际应用场景。

Amazon S3 是一种无服务器对象存储服务,是 AWS 的基础服务之一。SageMaker 利用 S3 处理各种用例,例如:

  • 用于存储训练数据集

  • 用于存储模型工件和训练输出

  • 存储异步端点和批量转换作业的推理输入和输出

Amazon S3 是一个高耐久性、可扩展且成本高效的存储解决方案。在访问存储在 S3 上的数据时,开发者可以选择将完整数据集从 S3 位置下载到 SageMaker 计算节点,或者通过流式传输数据。将大型数据集从 S3 下载到 SageMaker 计算节点会增加训练作业的启动时间。

Amazon EFS

Amazon Elastic File SystemEFS)是一个弹性文件系统服务。Amazon SageMaker 支持将训练数据集存储在 EFS 位置。在训练时,SageMaker 节点会挂载到 EFS 位置并直接访问训练数据集。在这种情况下,节点访问数据时不需要数据迁移,通常可以减少训练作业的启动时间。EFS 还允许多个节点持久化并无缝共享数据(因为这是一个共享系统)。当缓存或系统状态需要在训练节点之间共享时,这非常有用。

Amazon FSx for Lustre

Amazon FSx for Lustre 是专门为低延迟、高性能场景设计的共享文件系统服务。Amazon FSx 会自动从 S3 源复制数据,并使其在 SageMaker 计算节点上可用。Amazon FSx 具有与 EFS 相似的优势——即减少训练作业的启动时间,并提供共享文件系统。

编排服务

编排服务允许你将基于 SageMaker 的工作负载与其他 IT 生态系统进行集成。

AWS Step Functions

AWS Step Functions 是一个无服务器工作流服务,允许你编排业务流程和与其他 AWS 服务的交互。通过 Step Functions,你可以轻松地将单独的步骤组合成可重用和可部署的工作流。它支持可视化设计、分支和条件逻辑。

Step Functions 提供与 SageMaker 资源的原生集成。在需要通过多个服务协调复杂的机器学习流程时,Step Functions 非常有用。AWS 为开发者提供了 AWS Step Functions Data Science Python SDK 来开发、测试和执行此类流程。

Amazon API Gateway

Amazon API Gateway 是一个完全托管的 API 管理服务,用于开发、监控和管理 API。API Gateway 支持多个功能,在开发高度可扩展和安全的机器学习推理 API 时非常有用:

  • 认证和授权机制

  • 请求缓存、速率限制和流量控制

  • 防火墙功能

  • 请求头和负载转换

API Gateway 允许你将 SageMaker 实时端点与外部流量隔离,并提供额外的安全层。它还允许你为终端用户提供统一的 API,而无需暴露 SageMaker 运行时 API 的具体细节。

安全服务

强大的安全控制对于任何机器学习工作负载都是必不可少的,特别是在处理私密和敏感数据时。虽然本书并不专注于安全性,但了解 AWS 上权限和数据加密的基本概念非常重要。

AWS IAM

AWS 身份与访问管理 (IAM) 允许客户管理对 AWS 服务和资源的访问。在 SageMaker 的情况下,IAM 具有双重功能:

  • IAM 角色和策略定义了 SageMaker 任务可以访问和管理的 AWS 资源——例如,使用假定的 IAM 角色的训练任务是否可以访问 S3 上给定的数据集。

  • IAM 角色和策略定义了哪些主体(用户或服务)可以访问和管理 SageMaker 资源。例如,它定义了给定用户是否可以使用特定集群配置安排 SageMaker 训练任务。

审查 IAM 超出了本书的范围,但您需要了解它。在使用 SageMaker 时,设置 IAM 权限和角色是必要的前提。

Amazon VPC

Amazon 虚拟私有云 (VPC) 是一项服务,允许您在逻辑上隔离的私有网络中运行云工作负载。这种网络级别的隔离为谁可以访问您的工作负载提供了额外的安全性和控制。SageMaker 允许您在专用的 VPC 内运行训练和推理工作负载,这样您就可以控制从 SageMaker 资源进出的流量。

AWS KMS

AWS 密钥管理服务 (KMS) 用于加密底层数据。它还管理在加密和解密数据时对加密密钥的访问。在 SageMaker 中,KMS 主要用于加密训练数据、模型工件以及 SageMaker 集群中的底层磁盘。KMS 与所有可用的存储解决方案(如 S3、EFS 和 EBS(底层磁盘卷))集成。

监控服务

AWS 提供了专门的服务,用于监控其他 AWS 资源的管理、审计和执行。

Amazon CloudWatch

CloudWatch 提供监控和可观察性功能。在 SageMaker 的上下文中,它主要用于两个目的:

  • 用于存储和管理来自 SageMaker 资源(如端点或训练任务)的日志。默认情况下,SageMaker 会将 stdout/stderr 日志发送到 CloudWatch。

  • 用于存储时间序列指标。SageMaker 默认提供了若干指标(例如,对于实时端点,它会流式传输延迟和调用指标)。不过,开发人员也可以实现自定义指标。

Amazon CloudTrail

CloudTrail 捕获与管理任何 AWS 资源(包括 SageMaker 资源)相关的所有活动(如 API 调用)。通常,CloudTrail 用于治理和审计目的,但它也可以用来构建事件驱动的工作流。例如,开发人员可以使用它来监控资源创建或更新请求,并通过编程方式响应特定事件。

总结

本章开始时,我们提供了深度学习领域及其挑战的概述,以及 Amazon SageMaker 服务及其对深度学习工作负载的价值主张。然后,我们回顾了 SageMaker 的核心功能:托管训练和托管推理。我们研究了 SageMaker 训练作业和实时推理端点的生命周期。代码片段展示了如何使用 Python SDK 编程方式配置和提供 SageMaker 资源。我们还介绍了其他相关的 AWS 服务,因为在本书的其余部分中,我们将频繁使用它们。这将帮助我们更好地理解它们的使用和功能。

在下一章,我们将更深入地探讨任何 SageMaker 工作负载的基础构建块:运行时环境(特别是支持的深度学习框架)和容器。SageMaker 提供了多个流行的预配置运行时环境和容器,但它也允许通过其“BYO 容器”功能来完全自定义这些环境。我们将学习何时选择这些选项之一以及如何使用它们。

第二章:SageMaker 上的深度学习框架和容器

Amazon SageMaker 支持许多流行的机器学习(ML)和深度学习(DL)框架。SageMaker 中的框架支持通过使用预构建的 Docker 容器来实现推理和训练任务。预构建的 SageMaker 容器提供了大量的功能,并允许你在最少的编码工作下实现广泛的用例。也有一些实际场景,需要为训练和/或推理任务提供定制的运行时环境。为了解决这些情况,SageMaker 提供了灵活的自带容器BYO)功能。

在本章中,我们将回顾支持的主要深度学习框架及其对应的容器镜像。然后,我们将重点关注两个最流行的深度学习框架——TensorFlow 和 PyTorch,并学习如何在 Amazon SageMaker 中使用它们。此外,我们还将回顾一个更高层次的、用于自然语言处理任务的最先进框架 Hugging Face,以及它在 Amazon SageMaker 上的实现。

接着,我们将理解如何根据你的用例需求使用和扩展预构建的 SageMaker 容器,以及了解 SageMaker SDK 和工具包,这些工具简化了编写与 Amazon SageMaker 兼容的训练和推理脚本的过程。

在后续章节中,我们将深入探讨如何决定是使用预构建的 SageMaker 容器还是 BYO 容器。接着,我们将开发一个与 SageMaker 兼容的 BYO 容器。

以下部分将涵盖这些主题:

  • 探索 SageMaker 上的深度学习框架

  • 使用 SageMaker DL 容器

  • 开发 BYO 容器

到本章结束时,你将能够根据具体问题要求和选定的深度学习框架,决定选择哪种容器策略。此外,你将理解与 Amazon SageMaker 兼容的训练和推理脚本开发的关键方面。

技术要求

使用 SageMaker DL 容器开发 BYO 容器部分,我们将提供操作代码示例,以帮助你培养实践技能。完整的代码示例可以在 github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter2/ 查看。

为了跟随这段代码,你需要以下内容:

探索 SageMaker 上的深度学习框架

在撰写本书时,Amazon SageMaker 支持以下框架,其中以星号标记的为深度学习框架:

  • scikit-learn

  • SparkML 服务

  • Chainer*

  • Apache MXNet*

  • Hugging Face*

  • PyTorch*

  • TensorFlow*

  • 强化学习容器 – 包括支持 TensorFlow 和 PyTorch 的容器

  • XGBoost

上述支持的框架列表可能会在未来发生变化。请务必查看官方 SageMaker 文档:docs.aws.amazon.com/sagemaker/latest/dg/frameworks.xhtml

本书将主要聚焦于两种最受欢迎的选择:TensorFlowPyTorch。这两者都是开源框架,拥有庞大且充满活力的社区。根据具体的使用场景或模型架构,某一个框架可能会略有优势。然而,可以认为这两个框架在功能和性能上是相当的。在许多实际场景中,选择 TensorFlow 或 PyTorch 通常是基于历史惯例或个人偏好。

本书中我们还将讨论另一个框架:Hugging Face。这是一个高层次框架,提供了用于 NLP 任务(如文本分类、翻译等)的 SOTA 模型、训练和推理设施。Hugging Face 是一组旨在简化构建 SOTA NLP 模型的多个库(transformers、datasets、tokenizers 和 accelerate)。在底层,Hugging Face 库使用 TensorFlow 和 PyTorch 的基础组件(统称为“后端”)来执行计算。用户可以根据特定的运行时需求选择使用哪一个后端。鉴于其流行性,Amazon SageMaker 最近为 Hugging Face 库提供了支持,提供了分别用于训练和推理任务的预构建容器。

容器源

SageMaker 深度学习容器的源代码可以在 github.com/aws/deep-learning-containers 的公开 GitHub 仓库中找到。在某些情况下,查看相关的 Dockerfile 可以帮助你了解预构建容器的运行时配置。容器镜像可以在 AWS 公开注册表中找到,地址是 github.com/aws/deep-learning-containers/blob/master/available_images.md

对于每个支持的框架,SageMaker 提供了分别用于训练和推理的容器。我们为这两个任务提供了不同的容器,考虑因素如下:

  • 训练和推理任务可能有不同的运行时要求。例如,您可能会选择在不同的计算平台上运行训练和推理任务。根据您的具体任务,这将导致容器中使用不同的加速器和性能优化调整。

  • 训练和推理任务需要不同的辅助脚本;例如,在推理任务中启动模型服务器。如果不将训练和推理容器分开,可能会导致容器体积过大和 API 变得复杂。

因此,我们总是根据具体任务明确标识所使用的容器。

针对 DL 容器,AWS 还定义了独立的基于 GPU 和基于 CPU 的容器。基于 GPU 的容器需要安装额外的加速器,以便能够在 GPU 设备上运行计算(如 CUDA 工具包)。

模型要求

在选择 SageMaker DL 容器时,始终考虑模型对计算资源的要求。对于大多数 SOTA(最先进的)模型,建议使用基于 GPU 的计算实例,以实现可接受的性能。根据这一点选择适合的 DL 容器。

TensorFlow 容器

TensorFlow 容器有两个主要版本:1.x(维护模式)和 2.x(最新版本)。Amazon SageMaker 支持这两个版本,并提供推理和训练容器。本书中的所有代码示例和一般评论都假设使用 TensorFlow v2.x。

AWS 会频繁更新支持的小版本 TensorFlow。最新支持的主版本是 2.10.0。

PyTorch 容器

Amazon SageMaker 提供了用于 PyTorch 的推理和训练容器。最新版本是 1.12.1。

Hugging Face 容器

AWS 提供了两种版本的 Hugging Face 容器:PyTorch 和 TensorFlow 后端。每个后端都有独立的训练和推理容器。

使用 SageMaker Python SDK

AWS 提供了一个方便的 Python SDK,它通过 Estimator、Model 和 Predictor 类简化了与支持的 DL 框架的交互。每个支持的框架都有一个单独的模块,包含各自类的实现。例如,下面是如何为 PyTorch 框架导入 Predict、Estimator 和 Model 类:

from sagemaker.pytorch.estimator import PyTorch
from sagemaker.pytorch.model import PyTorchModel, PyTorchPredictor

以下图表展示了 SageMaker Python SDK 的工作流程:

图 2.1 – SageMaker Python SDK 如何与镜像 URI 配合使用

图 2.1 – SageMaker Python SDK 如何与镜像 URI 配合使用

为了更好地理解,让我们做一个简单的示例,展示如何使用 SageMaker Python SDK 使用特定版本的 PyTorch 容器运行训练作业。有关视觉概述,请参见 图 2.1

  1. 首先,我们决定使用哪个框架,并导入相应的 Pytorch estimator 类:

    from sagemaker.pytorch.estimator import PyTorch
    

在实例化 PyTorch estimator 对象时,我们需要提供更多参数,包括框架版本和 Python 版本:

estimator = PyTorch(
    entry_point="training_script.py",
    framework_version="1.8",
    py_version="py3",
    role=role,
    instance_count=1,
    instance_type="ml.p2.xlarge"
)
  1. 执行此代码时,SageMaker Python SDK 会自动验证用户输入,包括框架版本和 Python 版本。如果请求的容器存在,SageMaker Python SDK 将检索相应的容器镜像 URI。如果没有符合请求参数的容器,SageMaker Python SDK 会抛出异常。

  2. fit() 调用过程中,正确的容器镜像 URI 会被提供给 SageMaker API,因此训练作业将运行在安装了 PyTorch v1.8 和 Python v3.7 的 SageMaker 容器中。由于我们请求的是基于 GPU 的实例,将使用安装了 CUDA 工具包的训练容器:

    estimator.fit()
    

使用自定义镜像

请注意,如果由于某些原因您更倾向于提供容器镜像的直接 URI,您可以通过 image_uri 参数来实现,该参数在 modelestimator 类中受支持。

现在,让我们深入了解 SageMaker DL 容器,从可用的 TensorFlow、PyTorch 和 Hugging Face 框架的预构建容器开始。

使用 SageMaker DL 容器

Amazon SageMaker 支持多种容器使用模式。此外,它为您提供了训练和推理工具包,以简化使用预构建容器和开发 BYO 容器的过程。

在本节中,我们将学习如何选择最有效的容器使用模式来满足您的用例需求,并如何使用可用的 SageMaker 工具包来实现它。

容器使用模式

Amazon SageMaker 提供了灵活性,允许您选择是否使用预构建容器的“原样”(称为 脚本模式)、BYO 容器,或修改预构建容器。

通常,方法的选择取决于特定的模型运行时需求、可用资源以及工程专业知识。在接下来的几个子章节中,我们将讨论何时选择一种方法而不是另一种方法。

脚本模式

在脚本模式下,您需要定义要使用的预构建容器,然后提供一个或多个脚本来实现您的训练或推理逻辑。此外,您还可以提供任何其他依赖项(无论是专有的还是公共的),这些依赖项将被导出到容器中。

脚本模式下的训练和推理容器都预装了工具包,提供了常见功能,如将数据下载到容器和模型工件、启动作业等。我们将在本章稍后详细介绍 SageMaker 推理工具包训练工具包

脚本模式适用于以下场景:

  • 预构建容器满足您的运行时需求,或者您可以安装任何依赖项,而无需重建容器

  • 您希望最小化开发和测试容器所花费的时间,或者您没有所需的专业知识来进行此操作

在接下来的章节中,我们将回顾如何准备您的第一个训练和推理脚本,并以脚本模式在 SageMaker 上运行它们。

修改预构建容器

使用 SageMaker 预构建容器的另一种方式是对其进行修改。在这种情况下,您将使用其中一个预构建容器作为自定义容器的基础镜像。

在以下场景中,修改预构建容器可能会带来好处:

  • 你需要添加额外的依赖项(例如,必须从源代码编译的依赖项)或重新配置运行时环境

  • 你希望最小化容器的开发和测试工作,并在大部分情况下依赖 AWS 测试过的基础容器的功能

请注意,当你扩展预构建容器时,你需要负责以下几个方面:

  • 创建包含运行时环境实现的 Dockerfile

  • 在容器注册中心(例如Amazon Elastic Container RegistryECR)或私有 Docker 注册中心)中构建并存储容器

在本章稍后的部分,我们将看到如何扩展一个预构建的 PyTorch 容器来执行训练任务的示例。

自带容器(BYO containers)

你可能需要创建自定义容器的许多场景,包括以下几种:

  • 你有独特的运行时需求,这些需求无法通过扩展预构建容器来解决

  • 你希望为特定的硬件平台编译框架和库

  • 你使用的深度学习框架是 SageMaker 本地不支持的(例如,JAX)

构建一个与 SageMaker 推理和训练资源兼容的自定义容器需要开发工作、对 Docker 容器的理解以及特定的 SageMaker 要求。因此,通常建议首先考虑脚本模式或扩展一个预构建容器,只有在前两者无法满足你的具体用例时,才选择使用自带容器(BYO)。

SageMaker 工具包

为了简化与 Amazon SageMaker 兼容的自定义脚本和容器的开发,AWS 创建了用于训练和推理任务的 Python 工具包。

工具包提供以下好处:

  • 建立一致的运行时环境和存储代码资产的目录位置

  • ENTRYPOINT 脚本,用于容器启动时执行任务

理解这些工具包有助于简化并加速与 SageMaker 兼容容器的开发,因此让我们详细回顾一下它们。

训练工具包

SageMaker 训练工具包有几个关键功能:

  • 它建立了一个一致的运行时环境,设置了环境变量和一个目录结构来存储模型训练的输入和输出结果:

图 2.2 – SageMaker 兼容容器中的目录结构

图 2.2 – SageMaker 兼容容器中的目录结构

训练工具包会在训练容器中设置以下目录:

  • /opt/ml/input/config 目录,存放模型超参数和用于分布式训练的网络布局(以 JSON 文件形式)。

  • 使用 S3 作为数据存储时,/opt/ml/input/data 目录存放输入数据。

  • /opt/ml/code/ 目录,包含用于运行训练任务的代码资产。

  • /opt/ml/model/ 目录,包含生成的模型;SageMaker 在训练完成后会自动将其复制到 Amazon S3。

  • 它执行入口脚本并处理成功和失败状态。如果训练任务失败,输出将存储在/opt/ml/output/failure。对于成功的执行,工具包会将输出写入/opt/ml/success目录。

默认情况下,所有预构建的训练容器都已安装训练工具包。如果你希望使用它,你需要通过运行以下命令将其安装到容器中:

RUN pip install sagemaker-training

此外,你需要将所有代码依赖项复制到容器中,并在主训练脚本中定义一个特殊的环境变量,如下所示:

COPY train_scipt.py /opt/ml/code/train_script.py
ENV SAGEMAKER_PROGRAM train_scipt.py

训练工具包可在 PyPI(pypi.org)包和 SageMaker 的 GitHub 仓库(github.com/aws/sagemaker-training-toolkit)中找到。

推理工具包

推理工具包实现了一个与 SageMaker 推理服务兼容的模型服务堆栈。它与开源的多模型服务器MMS)一起,用于服务模型。它具有以下关键功能:

  • 用于建立运行时环境,例如存储推理输入和输出工件的目录和环境变量。目录结构遵循训练容器的布局。

  • 实现一个处理程序服务,该服务由模型服务器调用,将模型加载到内存中,并处理模型输入和输出。

  • 实现默认的序列化器和反序列化器来处理推理请求。

推理工具包可在 PyPi(pypi.org)包和 GitHub 仓库(github.com/aws/sagemaker-inference-toolkit)中找到。

开发脚本模式

现在我们对 SageMaker 的容器生态系统有了理解,接下来让我们实现几个学习项目来培养实际技能。在这个第一个示例中,我们将使用 SageMaker 脚本模式来训练我们的自定义 NLP 模型并部署它进行推理。

问题概述

在这个示例中,我们将学习如何使用 Hugging Face 框架开发训练和推理脚本。我们将利用 SageMaker 为 Hugging Face(具有 PyTorch 后端)预构建的容器。

我们选择了解决一个典型的 NLP 任务:文本分类。我们将使用20 Newsgroups数据集,它汇集了约 20,000 篇新闻组文档,涵盖 20 个不同的新闻组(类别)。有多种模型架构可以解决这个任务。通常,当前的 SOTA 模型基于 Transformer 架构。像BERT及其各种衍生模型这样的自回归模型适合这个任务。我们将使用一种被称为迁移学习的概念,其中一个为某个任务预训练的模型可以通过最小修改用于新任务。

作为基准模型,我们将使用被称为DistilBERT的模型架构,它在多种任务中提供了高准确度,并且比其他模型(例如原始的 BERT 模型)要小得多。为了将该模型适应分类任务,我们需要添加一个分类层,在我们的训练过程中,该层将被训练以识别文章:

图 2.3 – 文本分类任务的模型架构

图 2.3 – 文本分类任务的模型架构

Hugging Face Transformers 库通过以下方式简化了模型选择和修改,以便进行微调:

  • 它提供了丰富的模型库,包含许多预训练模型和分词器

  • 它提供了一个简单的模型 API,用于修改基准模型,以便对特定任务进行微调

  • 它实现了推理管道,将数据预处理和实际推理结合在一起

本学习项目的完整源代码可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter2/1_Using_SageMaker_Script_Mode.ipynb找到。

开发训练脚本

在运行 SageMaker 训练任务时,我们需要提供一个训练脚本。此外,我们还可以提供任何其他依赖项。我们还可以通过requirements.txt文件在预构建的容器上安装或修改已安装的 Python 包。

在这个例子中,我们将使用 Hugging Face 框架的一个新特性,通过 Hugging Face Trainer API 对多类别分类器进行微调。让我们确保训练容器中已安装较新的 Hugging Face Transformer 库。为此,我们创建requirements.txt文件并指定最小兼容版本。稍后,我们将把这个文件提供给我们的 SageMaker 训练任务:

transformers >= 4.10

接下来,我们需要开发训练脚本。让我们回顾一下其中的一些关键组件。

在训练时,SageMaker 通过调用user_training_script --arg1 value1 --arg2 value2 ...来启动训练。在这里,arg1..N是训练超参数和用户提供的其他杂项参数,这些参数作为训练任务配置的一部分提供。为了在我们的脚本中正确启动训练过程,我们需要在脚本中包含main guard

  1. 为了正确捕获参数,训练脚本需要能够解析命令行参数。我们使用 Python 的argparse库来实现这一点:

    if __name__ == "__main__":
         parser = argparse.ArgumentParser()
         parser.add_argument("--epochs", type=int, default=1)
         parser.add_argument("--per-device-train-batch-size", type=int, default=16)
         parser.add_argument("--per-device-eval-batch-size", type=int, default=64)
         parser.add_argument("--warmup-steps", type=int, default=100)
         parser.add_argument("--logging-steps", type=float, default=100)
         parser.add_argument("--weight-decay", type=float, default=0.01)
         args, _ = parser.parse_known_args()
         train(args)
    
  2. train()方法负责执行端到端的训练任务。它包含以下组件:

    • 调用_get_tokenized_dataset,使用来自 Hugging Face 库的预训练 DistilBERT 分词器加载并分词数据集。

    • 从 Hugging Face 模型库加载并配置 DistilBERT 模型。请注意,我们更新了分类任务的默认配置,以调整我们选择的类别数量。

    • 配置 Hugging Face Trainer 并开始训练过程。

    • 训练完成后,我们保存训练好的模型:

      def train(args):
          train_enc_dataset, test_enc_dataset = _get_tokenized_data()
          training_args = TrainingArguments(
              output_dir=os.getenv(
                  "SM_OUTPUT_DIR", "./"
              ),  # output directory, if runtime is not
              num_train_epochs=args.epochs,
              per_device_train_batch_size=args.per_device_train_batch_size,
              per_device_eval_batch_size=args.per_device_eval_batch_size,
              warmup_steps=args.warmup_steps,
              weight_decay=args.weight_decay,
              logging_steps=args.logging_steps,
          )
          config = DistilBertConfig()
          config.num_labels = NUM_LABELS
          model = DistilBertForSequenceClassification.from_pretrained(
              MODEL_NAME, config=config
          )
          trainer = Trainer(
              model=model,  # model to be trained
              args=training_args,  # training arguments, defined above
              train_dataset=train_enc_dataset,  # training dataset
              eval_dataset=test_enc_dataset,  # evaluation dataset
          )
          trainer.train()
          model.save_pretrained(os.environ["SM_MODEL_DIR"])
      

到目前为止,我们的脚本已经涵盖了几个关键方面:处理配置设置和模型超参数、加载预训练模型,并使用 Hugging Face Trainer API 启动训练。

启动训练作业

一旦我们准备好了训练脚本和依赖项,就可以开始训练并通过 SageMaker Python SDK 调度训练作业。我们从导入 Hugging Face Estimator 对象并获取训练作业的 IAM 执行角色开始:

from sagemaker.huggingface.estimator import HuggingFace
from sagemaker import get_execution_role
role=get_execution_role()

接下来,我们需要定义模型和训练过程的超参数。这些变量将在训练时传递给我们的脚本:

hyperparameters = {
    "epochs":1,
    "per-device-train-batch-size":16, 
    "per-device-eval-batch-size":64,
    "warmup-steps":100,
    "logging-steps":100,
    "weight-decay":0.01    
}
estimator = HuggingFace(
    py_version="py36",
    entry_point="train.py",
    source_dir="1_sources",
    pytorch_version="1.7.1",
    transformers_version="4.6.1",
    hyperparameters=hyperparameters,
    instance_type="ml.p2.xlarge",
    instance_count=1,
    role=role
)
estimator.fit({
    "train":train_dataset_uri,
    "test":test_dataset_uri
})

之后,训练作业将被调度并执行。完成训练需要大约 10 到 15 分钟,训练后的模型和其他输出成果将被添加到 Amazon S3。

为脚本模式开发推理脚本

现在我们已经有了训练好的模型,接下来将其部署为 SageMaker 实时端点。我们将使用预构建的 SageMaker Hugging Face 容器,并只提供推理脚本。推理请求将由 AWS MMS 处理,并暴露 HTTP 端点。

使用预构建推理容器时,SageMaker 会自动识别我们的推理脚本。根据 SageMaker 的约定,推理脚本必须包含以下方法:

  • model_fn(model_dir) 在容器启动时执行,用于将模型加载到内存中。该方法以模型目录作为输入参数。你可以使用 model_fn() 来初始化推理管道中的其他组件,比如在我们的例子中是分词器。需要注意的是,Hugging Face Transformers 提供了一个方便的 Pipeline API,允许我们将数据预处理(在我们这里是文本分词)和实际推理结合成一个对象。因此,我们返回的是一个推理管道,而不是加载的模型:

    MODEL_NAME = "distilbert-base-uncased"
    NUM_LABELS = 6 # number of categories
    MAX_LENGTH = 512 # max number of tokens model can handle
    def model_fn(model_dir):
        device_id = 0 if torch.cuda.is_available() else -1
        tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_NAME)
        config = DistilBertConfig()
        config.num_labels = NUM_LABELS
        model = DistilBertForSequenceClassification.from_pretrained(
            model_dir, config=config
        )
        inference_pipeline = pipeline(
            model=model,
            task="text-classification",
            tokenizer=tokenizer,
            framework="pt",
            device=device_id,
            max_length=MAX_LENGTH,
            truncation=True
        )
        return inference_pipeline
    
  • transform_fn(inference_pipeline, data, content_type, accept_type) 负责执行实际的推理操作。由于我们通过 HTTP 与最终客户端进行通信,因此我们还需要进行负载反序列化和响应序列化。在我们的示例中,我们预期使用 JSON 负载并返回 JSON 负载;然而,依据需求,这可以扩展为其他格式(例如 CSV 和 Protobuf)。

    def transform_fn(inference_pipeline, data, content_type, accept_type):
        # Deserialize payload
        if "json" in content_type:
            deser_data = json.loads(data)
        else:
            raise NotImplemented("Only 'application/json' content type is implemented.")
    
        # Run inference
        predictions = inference_pipeline(deser_data)
    
        # Serialize response
        if "json" in accept_type:
            return json.dumps(predictions)
        else:
            raise NotImplemented("Only 'application/json' accept type is implemented.")
    

有时候,将反序列化、推理和序列化结合到一个方法中可能会不方便。作为替代,SageMaker 支持更细粒度的 API:

  • input_fn(request_body, request_content_type) 执行反序列化操作。

  • predict_fn(deser_input, model) 执行预测操作

  • output_fn(prediction, response_content_type) 执行预测结果的序列化

请注意,transform_fn()方法与input_fn()predict_fn()output_fn()方法是互斥的。

部署文本分类终端节点

现在我们准备好部署并测试我们的 Newsgroup 分类终端节点了。我们可以使用estimator.create_model()方法来配置我们的模型部署参数,特别是以下内容:

  1. 定义推理脚本和其他依赖项,这些内容将由 SageMaker 上传至终端节点。

  2. 确定推理容器。如果提供transformers_versionpytorch_versionpy_version参数,SageMaker 会自动找到一个适当的预构建推理容器(如果存在)。另外,您可以提供image_uri来直接指定希望使用的容器镜像:

    from sagemaker.huggingface.estimator import HuggingFaceModel
    model = estimator.create_model(role=role, 
                                   entry_point="inference.py", 
                                   source_dir="1_sources",
                                   py_version="py36",
                                   transformers_version="4.6.1",
                                   pytorch_version="1.7.1"
                                  )
    
  3. 接下来,我们定义终端节点的参数,例如其背后的实例数量和类型。model.deploy()方法启动推理部署(通常需要几分钟时间),并返回一个Predictor对象以运行推理请求:

    predictor = model.deploy(
        initial_instance_count=1,
        instance_type="ml.m5.xlarge"
    )
    

接下来,让我们探讨如何扩展预构建的深度学习容器。

扩展预构建容器

我们将重用脚本模式示例中的代码资产。然而,与之前的容器不同,我们将修改我们的运行时环境,并从 GitHub 主分支安装最新稳定版的 Hugging Face Transformer。此修改将在我们的自定义容器镜像中实现。

首先,我们需要确定将使用哪个基础镜像。AWS 已在github.com/aws/deep-learning-containers/blob/master/available_images.md发布了所有可用的深度学习容器。

由于我们计划重新从头安装 HuggingFace Transformer 库,因此可以选择 PyTorch 基础镜像。在撰写时,最新的 PyTorch SageMaker 容器是763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.9.0-gpu-py38-cu111-ubuntu20.04。请注意,此容器 URI 适用于 AWS East-1 区域,其他 AWS 区域会有所不同。请查阅前述参考的 AWS 文章以了解适用于您区域的正确 URI。

要构建一个新的容器,我们需要执行以下步骤:

  • 创建一个包含运行时指令的 Dockerfile。

  • 在本地构建容器镜像。

  • 将新的容器镜像推送到容器注册表。在本示例中,我们将使用 ECR 作为容器注册表:AWS 提供的托管服务,它与 SageMaker 生态系统深度集成。

首先,让我们为我们的扩展容器创建一个 Dockerfile。

为我们的扩展容器开发 Dockerfile。

为了扩展预构建的 SageMaker 容器,我们至少需要以下组件:

  • 一个用于作为基础的 SageMaker PyTorch 镜像。

  • 安装所需的依赖项,如最新版本的 PyTorch 和来自最新 Git 主分支的 Hugging Face Transformers。

  • 将我们之前示例中的训练脚本复制到容器中。

  • 定义 SAGEMAKER_SUBMIT_DIRECTORYSAGEMAKER_PROGRAM 环境变量,以便 SageMaker 知道容器启动时应该执行哪个训练脚本:

    FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.9.0-gpu-py38-cu111-ubuntu20.04
    RUN pip3 install git+https://github.com/huggingface/transformers 
    ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code 
    ENV SAGEMAKER_PROGRAM train.py 
    COPY 1_sources/train.py $SAGEMAKER_SUBMIT_DIRECTORY/$SAGEMAKER_PROGRAM
    

现在我们准备好构建并将此容器镜像推送到 ECR。你可以在本章的代码库中找到执行此操作的 bash 脚本。

安排训练作业

一旦我们在 ECR 中拥有扩展的 PyTorch 容器,就可以执行一个 SageMaker 训练作业了。训练作业配置将与脚本模式示例类似,唯一的不同点是:我们将使用一个通用的 SageMaker Estimator 对象,而不是 HuggingFaceEstimator 对象,后者可以让我们使用自定义镜像。请注意,您需要更新 image_uri 参数,指向您 ECR 实例中的镜像 URI。你可以通过访问 AWS 控制台中的 ECR 服务来找到它,并在那找到扩展容器:

from sagemaker.estimator import Estimator
estimator = Estimator(
    image_uri="<UPDATE WITH YOUR IMAGE URI FROM ECR>",
    hyperparameters=hyperparameters,
    instance_type="ml.p2.xlarge",
    instance_count=1,
    role=role
)
estimator.fit({
    "train":train_dataset_uri,
    "test":test_dataset_uri
})

完成训练作业后,我们应当期望与脚本模式示例中展示的训练结果类似。

开发一个用于推理的 BYO 容器

在本节中,我们将学习如何使用官方 TensorFlow 镜像构建一个与 SageMaker 兼容的推理容器,准备推理脚本和模型服务器,并将其部署到 SageMaker Hosting 上进行推理。

问题概述

我们将开发一个与 SageMaker 兼容的推理容器。我们将使用最新的官方 TensorFlow 容器作为基础镜像,并使用 AWS MMS 作为模型服务器。请注意,MMS 是许多机器学习模型服务选项之一,SageMaker 对模型服务器没有其他限制,唯一要求是它应该在端口8080上提供模型服务。

开发服务容器

当将服务容器部署到端点时,SageMaker 会运行以下命令:

docker run <YOUR BYO IMAGE> serve

为了遵守这一要求,建议在 Dockerfile 中使用 ENTRYPOINT 指令的执行格式。

让我们回顾一下我们的 BYO Dockerfile:

  • 我们使用最新的 TensorFlow 容器作为基础

  • 我们安装通用和 SageMaker 特定的依赖项

  • 我们将模型服务脚本复制到容器中

  • 我们指定 ENTRYPOINT 和 CMD 指令,以遵守 SageMaker 的要求

现在,让我们开始行动:

  1. 使用最新的官方 TensorFlow 容器:

    FROM tensorflow/tensorflow:latest
    
  2. 安装 Java,因为 MMS 和其他常见依赖项都需要它。

  3. 将入口脚本复制到镜像中:

    COPY 3_sources/src/dockerd_entrypoint.py /usr/local/bin/dockerd-entrypoint.py
    RUN chmod +x /usr/local/bin/dockerd-entrypoint.py
    
  4. 将默认的自定义服务文件复制到容器中,以处理传入的数据和推理请求:

    COPY 3_sources/src/model_handler.py /opt/ml/model/model_handler.py
    COPY 3_sources/src/keras_model_loader.py /opt/ml/model/keras_model_loader.py
    
  5. 定义入口脚本及其默认参数:

    ENTRYPOINT ["python3", "/usr/local/bin/dockerd-entrypoint.py"]
    CMD ["serve"]
    

在这个例子中,我们并不打算详细介绍 MMS 以及推理脚本的开发。然而,值得强调一些关键的脚本方面:

  • dockerd_entrypoint.py 是一个可执行文件,当传递 serve 参数时,它会启动 MMS 服务器。

  • model_handler.py 实现了模型加载和服务逻辑。请注意,handle() 方法检查模型是否已经加载到内存中。如果没有,它会先加载模型到内存,然后继续处理服务请求,包括以下内容:

    • 反序列化请求负载

    • 运行预测

    • 序列化预测

部署 SageMaker 端点

要安排推理端点的部署,我们使用来自 SageMaker Python SDK 的通用 Model 类。请注意,由于我们从公共模型库下载了模型,因此无需提供 model_data 参数(因此其值为 None):

from sagemaker import Model
mms_model = Model(
    image_uri=image_uri,
    model_data=None,
    role=role,
    name=model_name,
    sagemaker_session=session
)
mms_model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge", 
    endpoint_name=endpoint_name
)

完全部署端点并启动模型服务器可能需要几分钟。一旦准备就绪,我们可以使用 boto3.sagemaker-runtime 客户端调用端点,该客户端允许你构建 HTTP 请求并将推理负载(或者在我们的案例中是图像)发送到特定的 SageMaker 端点:

import boto3
client = boto3.client('sagemaker-runtime')
accept_type = "application/json"
content_type = 'image/jpeg'
headers = {'content-type': content_type}
payload = open(test_image, 'rb')
response = client.invoke_endpoint(
    EndpointName=endpoint_name,
    Body=payload,
    ContentType=content_type,
    Accept = accept_type
)
most_likely_label = response['Body'].read()
print(most_likely_label)

这段代码很可能会根据模型预测返回图像中的对象。

总结

在本章中,我们回顾了 SageMaker 如何使用 Docker 容器支持机器学习和深度学习框架。阅读完本章后,你现在应该知道如何根据特定的使用案例需求选择最合适的深度学习容器使用模式。我们了解了 SageMaker 工具包,这简化了开发与 SageMaker 兼容的容器。在后续部分,你将获得如何为 Amazon SageMaker 开发自定义容器和脚本进行训练和推理任务的实用知识。

在下一章中,我们将学习 SageMaker 开发环境,以及如何高效地开发和排查你的深度学习代码。此外,我们还将了解 SageMaker 开发环境提供的深度学习特定工具和接口,这些工具可以简化深度学习模型的构建、部署和监控。

第三章:管理 SageMaker 开发环境

在前面的章节中,我们学习了 Amazon SageMaker 的基本组件和功能。到现在为止,你已经知道如何在 SageMaker 上构建和部署第一个简单的模型。然而,在许多更复杂的情况下,你需要在将代码部署到 SageMaker 管理的训练或托管集群之前,编写、分析和测试你的深度学习(DL)代码。能够在本地执行这一操作,同时模拟 SageMaker 运行时,将缩短开发周期,并避免为开发提供 SageMaker 资源所带来的不必要成本。

在本章中,我们将探讨如何组织你的开发环境,以便有效地为 SageMaker 开发和测试你的 DL 模型。本章包括选择 IDE 软件进行开发和测试的注意事项,以及如何在本地机器上模拟 SageMaker 运行时环境。我们还将概述用于管理 SageMaker 资源的可用 SDK 和 API。

以下主题将在后续章节中进行讲解:

  • 为 SageMaker 选择开发环境

  • 在本地调试 SageMaker 代码

阅读完本章后,你将能够根据特定的使用案例要求,设置与 SageMaker 兼容的高效开发环境。

技术要求

在本章中,你可以通过代码示例进行演练,从而培养实际技能。完整的代码示例可以在这里找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter3/。为了跟随这些代码,你需要具备以下条件:

  • 拥有 AWS 账户,并且是具有管理 Amazon SageMaker 资源权限的 IAM 用户。

  • 在本地机器上安装 Docker 和 Docker Compose。如果你的开发环境中有 GPU 设备,你还需要安装 nvidia-docker (github.com/NVIDIA/nvidia-docker)。

  • 安装 Conda (docs.conda.io/en/latest/)。

为 SageMaker 选择开发环境

开发环境和 IDE 的选择通常由个人偏好或公司政策驱动。由于 SageMaker 是一个云平台,它不会限制你使用任何你选择的 IDE。你可以在本地机器或云机器(例如 Amazon EC2)上运行 IDE。SageMaker 还提供了一套 SDK 和软件包,用于模拟 SageMaker 运行时环境,这样你可以在将任何代码部署到云端之前,先在本地模拟环境中测试代码。

随着数据科学,特别是机器学习的进步,一种新的开发运行时环境应运而生——交互式笔记本,即 Jupyter NotebooksJupyterLab(Jupyter Notebooks 的下一代,具有更多开发能力,如代码调试)。虽然它们并未完全取代传统 IDE,但笔记本因其能够探索和可视化数据、开发并与他人共享代码的功能而变得流行。

SageMaker 提供了几种托管的笔记本环境:

  • SageMaker Studio 服务 —— 一个专有的无服务器笔记本 IDE,用于机器学习开发

  • SageMaker 笔记本实例 —— 一种托管的 Jupyter Notebook/JupyterLab 环境

三种选择——传统的 IDE、SageMaker 笔记本实例和 SageMaker Studio——都有各自的优点,并且在特定场景下可能是最优的选择。在接下来的部分,我们将详细审视这些 IDE 选项,并讨论它们在深度学习开发中的优缺点。

为 SageMaker 设置本地环境

在本地进行初步开发有许多好处,具体包括以下几点:

  • 在本地进行开发时,您不会产生任何运行成本。

  • 您可以选择自己喜欢的 IDE,从而提高开发周期的效率。

然而,本地开发运行时也有一定的限制。例如,您无法在不同的硬件设备上测试和分析代码。获取最新的为深度学习工作负载设计的 GPU 设备可能不切实际且成本高昂。因此,在许多情况下,您会使用 CPU 设备进行深度学习代码的初步开发和测试,以解决初期问题,然后在访问目标 GPU 设备的云实例上进行最终的代码分析和调优。

SageMaker 提供了多个 SDK,允许在本地环境与 AWS 云之间进行集成。让我们通过一个实际例子来演示如何配置本地环境以使用远程 SageMaker 资源。

配置 Python 环境

我们通过设置并配置一个与 AWS 集成的 Python 环境来开始配置。建议使用 Conda 环境管理软件来隔离您的 SageMaker 本地环境:

  1. 您可以通过使用适当的安装方法(取决于您的本地操作系统)在本地机器上安装 Conda。安装完成后,您可以通过在终端窗口中运行以下命令来创建一个新的 Python 环境:

    conda create -n sagemaker python=3.9
    

请注意,我们在此环境中明确指定使用的 Python 解释器版本。

  1. 接下来,我们切换到创建环境并安装 AWS 和 SageMaker SDK:

    conda activate sagemaker
    pip install boto3 awscli sagemaker
    

让我们回顾一下刚刚安装的 SDK:

  • awscli 是一个 AWS CLI 工具包,允许您以编程方式与任何 AWS 服务进行交互。它还提供了一个机制,用于在本地存储和使用 AWS 凭证。

  • boto3 是一个 Python SDK,用于管理你的 AWS 资源。它使用 AWS CLI 工具包建立的凭证,通过加密签名任何管理请求,从而在 AWS 中进行身份验证。

  • sagemaker – 你应该已经熟悉这个 Python SDK,因为在本书的前几章中,我们使用它与 SageMaker 资源进行交互,如训练作业或推理端点。与 boto3 不同,SageMaker SDK 抽象了许多底层资源管理的方面,通常建议在你需要编程管理 SageMaker 工作负载时使用。

  1. 在继续之前,我们需要先配置 AWS 凭证。为此,你需要在终端中运行以下命令并提供你的 AWS 访问密钥和秘密密钥:

    aws configure
    

你可以在这里阅读有关如何设置 AWS 凭证的详细信息:docs.aws.amazon.com/cli/latest/userguide/cli-configure-quickstart.xhtml

配置 Jupyter 环境

一旦我们配置好基本的 Python 环境并设置好 AWS 凭证,我们就可以启动 Jupyter 服务器了。在这个例子中,我们将使用 JupyterLab 环境。然而,你也可以根据自己的需求配置 IDE,因为许多 IDE(如 PyCharm 和 Visual Studio Code)都支持通过插件或原生方式使用 Jupyter Notebook。此方法的额外好处是,你可以在同一 IDE 中轻松地在笔记本和训练、推理脚本之间切换:

  1. 要安装 JupyterLab 并创建内核,请在终端中运行以下命令:

    conda install -c conda-forge jupyterlabpython -m ipykernel install --user --name sagemaker
    
  2. 接下来,我们在本地机器上启动 JupyterLab 服务器:

    jupyter lab
    

你的 JupyterLab 服务器现在应该可以通过 http://localhost:8888 访问。

在 SageMaker 上运行模型训练

在 JupyterLab 实例中,我们运行一些测试以确保能够从本地机器连接和管理 SageMaker 资源:

  1. 本书 GitHub 仓库中的 chapter3 目录下包含完整的笔记本代码和训练脚本:

    import sagemaker, boto3
    from sagemaker import get_execution_role
    session = sagemaker.Session()
    account = boto3.client('sts').get_caller_identity().get('Account')
    role = f"arn:aws:iam::{account}:role/service-role/AmazonSageMaker-ExecutionRole-<YOUR_ROLE_ID>" 
    

SageMaker 执行角色

请注意,你需要手动定义执行角色。对于 SageMaker 管理的环境,如 SageMaker Studio 或 SageMaker Notebook 实例,你可以使用 get_execution_role() 方法来获取执行角色。

  1. 现在,我们可以像之前一样配置并启动 SageMaker 训练:

    from sagemaker.pytorch import PyTorch
    import os
    pytorch_estimator = PyTorch(
                            session=session,
                            entry_point=f'{os.getcwd()}/sources/cifar10.py',
                            role=role,
                            instance_type="ml.m4.xlarge",
                            instance_count=1,
                            job_name="test",
                            framework_version="1.9.0",
                            py_version="py38",
                            hyperparameters={
                                "epochs": 1,
                                "batch-size": 16
                                }
                            )
    pytorch_estimator.fit()
    
  2. 一旦训练作业完成,你可以查看本地的训练结果以及输出的工件存储位置:

    pytorch_estimator.latest_training_job.describe()
    

如你所见,拥有本地开发环境为你提供了选择首选 IDE 的灵活性,同时避免了为 SageMaker 管理的开发环境付费。与此同时,它要求你仔细管理开发环境,这需要特定的专业知识和投入的努力。另一个潜在挑战是团队成员之间开发环境的同步。

使用 SageMaker Notebook 实例

SageMaker 笔记本实例由一个运行在 EC2 实例上的 AWS Jupyter 环境管理。你可以从基于 CPU 和 GPU 的实例列表中选择实例类型。SageMaker 提供了多个预配置的 Jupyter 内核,支持 Python 运行时。它包括预配置的运行时,支持 PyTorch、TensorFlow、MXNet 和其他流行的深度学习和机器学习框架。你还可以自定义现有内核(例如,安装新软件包)或使用 Conda 环境管理创建完全自定义的内核。

由于 Jupyter 环境直接运行在 EC2 实例上,你可以在本地训练或推理时直接观察资源消耗(例如,通过监控 nvidia-smi 工具输出)。你还可以执行 Docker 操作,如构建自定义容器并使用 SageMaker 本地模式进行测试,我们将在本章的 本地调试 SageMaker 代码 部分中详细讨论。

在某些场景下,使用笔记本实例可能会带来好处,例如以下情况:

  • 你需要访问特定类型的硬件来测试和调试模型(例如,寻找最大训练吞吐量,而不遇到 OOM 问题)

  • 在将模型部署到远程环境之前,你希望在本地为特定的超参数组合和硬件基准化模型性能

笔记本实例的一个缺点是缺乏灵活性。如果硬件需求发生变化,你无法快速更改实例类型。这可能导致在任务组合中有不同资源需求时产生不必要的成本。

假设你需要在本地预处理训练数据,并在这些数据上调试训练脚本。通常,数据处理是一个 CPU 密集型过程,不需要 GPU 设备。然而,训练深度学习模型则需要 GPU 设备。因此,你必须为任务中最高硬件需求的部分提供一个实例。或者,你需要在任务之间存储工作并重新配置笔记本实例。

SageMaker 在一个名为 SageMaker Studio 笔记本的较新产品中解决了这一弹性不足的问题。我们来详细回顾一下。

使用 SageMaker Studio

SageMaker Studio是一个基于 Web 的界面,允许你与各种 SageMaker 功能进行交互,从可视化数据探索、模型库和模型训练,到代码开发和端点监控。SageMaker Studio 旨在通过提供一个单一的工作和协作环境,简化和优化机器学习开发的所有步骤。

SageMaker Studio 中有多种功能。我们来回顾一下与深度学习开发相关的两项具体功能:

  • Studio 笔记本允许快速访问不同的计算实例和运行时,无需离开 JupyterLab 应用程序

  • SageMaker JumpStart是一个预构建的解决方案和模型库,允许你通过几次点击部署你的深度学习解决方案。

接下来,让我们讨论这些功能和使用案例。

Studio 笔记本

Studio 笔记本提供了一个完全托管的 JupyterLab 环境,可以快速在不同的内核和计算实例之间切换。在切换过程中,你的工作会自动保存在共享文件系统中。共享文件系统具有高可用性,并根据需要无缝扩展。Studio 笔记本配备了一组预构建的内核,类似于笔记本实例,并且可以进一步自定义。你还可以为 Studio 笔记本创建一个完全自定义的内核镜像。

你可以从广泛的 EC2 实例、最新的 CPU 实例以及专用的 GPU 实例中选择用于训练和推理任务的计算实例。Studio 笔记本可以访问两种类型的 EC2 实例:

  • 快速实例,可以在 2 分钟内完成切换。

  • 常规实例,启动大约需要 5 分钟。请注意,这是一个大致的时间,可能会受到特定 AWS 区域资源可用性的影响。

协作功能

Studio 笔记本支持共享功能,允许你只需几次点击便能与团队成员共享代码、内核和实例配置。

SageMaker 笔记本内核在 Docker 镜像中运行。因此,存在若干限制:

  • 你不能在 Studio 笔记本中构建或运行容器。

  • Studio 笔记本不支持在部署到 SageMaker 之前调试容器的本地模式。

  • AWS 提供了Image Build CLI来绕过这个限制,允许用户在使用 Studio 笔记本时构建自定义容器。

在大多数场景下,Studio 笔记本将是运行你自己的 JupyterLab 在 EC2 实例上或使用 SageMaker 笔记本实例的一个方便且具有成本效益的替代方案。然而,你应该留意之前提到的 Studio 笔记本的限制,并评估这些限制是否会影响你的特定用例或使用模式。此外,Studio 笔记本是 SageMaker Studio 平台的一部分,提供了更多的额外功能,如可视化数据探索和处理、可视化模型监控、预构建解决方案、用于管理特征存储、模型构建管道、端点、实验等的 UI 便捷功能。

SageMaker JumpStart

SageMaker JumpStart 是一个预构建的端到端机器学习(ML)和深度学习(DL)解决方案库,提供可在 SageMaker 上一键部署的示例笔记本和模型。JumpStart 的解决方案和模型库庞大并持续增长。

JumpStart 解决方案专为特定行业用例设计,例如交易欺诈检测、文档理解和预测性维护。每个解决方案都包括多个集成组件,部署后可以立即供最终用户使用。请注意,您需要提供自己的数据集来训练 JumpStart 模型。

JumpStart 模型提供访问 SOTA 模型库。根据您的模型架构,您可以选择立即将该模型部署到推理、进行微调、从头训练或在自己的数据集上恢复增量训练。JumpStart 允许用户完全自定义用户操作,如定义训练集群的大小和实例类型、训练作业的超参数以及数据的位置。

模型库包括来自 TensorFlow Hub、PyTorch Hub 和 Hugging Face 的 CV 和 NLP 任务模型。

当您的业务问题可以通过使用通用解决方案和专有数据来解决时,SageMaker JumpStart 可以派上用场。JumpStart 还可以作为向 SageMaker 上的深度学习(DL)友好介绍,或者适合那些希望自己尝试深度学习的非技术用户。

在本节中,我们回顾了 SageMaker 可用的开发环境选项。所有三个选项都有其优缺点,具体选择主要取决于个人偏好和用例需求。通常,最好同时拥有本地环境和 SageMaker Studio 笔记本或笔记本实例。这种设置允许您在不支付任何云资源费用的情况下,本地开发、测试和进行初步调试。一旦您的代码在本地工作,您就可以轻松地在云硬件上运行相同的代码。Studio 笔记本尤其有用,因为它们允许您轻松切换不同的 CPU 和 GPU 运行时,而无需离开 Jupyter 笔记本,因此您可以实验训练配置(例如,调整批量大小或梯度累积)。

在下一节中,我们将重点介绍如何在将工作负载迁移到 SageMaker 云资源之前,高效地在本地调试 SageMaker 代码。

在本地调试 SageMaker 代码

为了简化本地代码开发和测试,SageMaker 支持 本地模式。此模式允许您在 SageMaker 容器中本地运行训练、推理或数据处理。当您希望在配置任何 SageMaker 资源之前先排查脚本问题时,这尤其有帮助。

所有 SageMaker 镜像以及自定义的 SageMaker 兼容镜像都支持本地模式。它作为 sagemaker Python SDK 的一部分实现。当您在本地模式下运行作业时,SageMaker SDK 会在后台创建一个包含作业参数的 Docker Compose YAML 文件,并在本地启动相关容器。配置 Docker 运行时环境的复杂性对用户进行了抽象。

本地模式支持 CPU 和 GPU 设备。你可以在本地模式下运行以下类型的 SageMaker 作业:

  • 训练作业

  • 实时端点

  • 处理作业

  • 批处理转换作业

本地模式的限制

在本地运行 SageMaker 作业时有一些限制:

  • 仅支持一个本地端点。

  • 不支持 GPU 的分布式本地训练。然而,你可以在 CPU 上运行分布式作业。

  • EFS 和 FSx for Lustre 不支持作为数据源。

  • 不支持Gzip压缩、管道模式或输入的清单文件。

在本地模式下运行训练和推理

让我们在本地模式下训练一个简单的模型,然后将推理端点本地部署。完整的笔记本代码和训练脚本位于书籍仓库的chapter3目录:

  1. 我们首先安装所有本地模式所需的依赖项:

    pip install 'sagemaker[local]' –upgrade
    
  2. 然后,我们配置 SageMaker 本地运行时。请注意,我们使用LocalSession类来让 SageMaker SDK 知道我们希望在本地配置资源:

    import boto3
    from sagemaker.local import LocalSession
    sagemaker_local_session = LocalSession()
    sagemaker_local_session.config = {'local': {'local_code': True}}
    account = boto3.client('sts').get_caller_identity().get('Account')
    role = f"arn:aws:iam::{account}:role/service-role/AmazonSageMaker-ExecutionRole-<YOUR_ROLE_ID>" 
    
  3. 在本笔记本中,我们打算使用来自 SageMaker ECR 仓库的公共 PyTorch 镜像。为此,我们需要存储凭证,以便 Docker 守护进程可以拉取镜像。在笔记本中运行以下命令(你也可以在终端窗口中运行,只需移除!):

    ! aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com
    
  4. 现在,我们需要决定是使用 GPU(如果可用)还是 CPU 设备(默认选择)。以下代码片段判断是否有可用的 CUDA 兼容设备("local_gpu"值),如果没有,则默认使用 CPU 设备("local"值):

    import subprocess
    instance_type = "local"
    try:
        if subprocess.call("nvidia-smi") == 0:
            instance_type = "local_gpu"
    except:
        print("GPU device with CUDA is not available")
    print("Instance type = " + instance_type)
    
  5. 一旦我们定义了使用的本地设备,我们就可以配置并运行 SageMaker 训练作业:

    from sagemaker.pytorch import PyTorch
    import os
    # Configure an MXNet Estimator (no training happens yet)
    pytorch_estimator = PyTorch(
                            session=sagemaker_local_session,
                            entry_point=f'{os.getcwd()}/sources/cifar10.py',
                            role=role,
                            instance_type=instance_type,
                            instance_count=1,
                            job_name="test",
                            framework_version="1.9.0",
                            py_version="py38",
                            hyperparameters={
                                "epochs": 1,
                                "batch-size": 16
                                }
                            )
    pytorch_estimator.fit()
    
  6. SageMaker Python SDK 会自动执行以下操作:

    • 从公共 ECR 仓库拉取适当的 PyTorch 镜像

    • 生成一个适当的docker-compose.yml文件,设置合适的卷挂载点以访问代码和训练数据

    • 使用train命令启动一个 Docker 容器

SageMaker 将输出 Docker Compose 命令及训练容器的 STDOUT/STDERR 到 Jupyter 单元格。

容器内调试代码

许多现代 IDE 支持调试在容器内运行的应用程序。例如,你可以在训练代码中设置断点。容器中的代码执行将停止,这样你就可以检查它是否正确执行。请查阅你的 IDE 文档,了解如何进行设置。

训练作业完成后,让我们看看如何将训练好的模型部署到本地实时端点。请注意,默认情况下,我们只训练单个 epoch,因此不要期望很好的结果!

  1. 你可以通过在估算器上运行deploy()方法,将推理容器本地部署:

    pytorch_estimator.deploy(initial_instance_count=1, instance_type=instance_type)
    
  2. 一旦端点部署完成,SageMaker SDK 会开始将模型服务器的输出发送到 Jupyter 单元格。你也可以在 Docker 客户端 UI 中或通过docker logs CONTAINER_ID终端命令观察容器日志。

  3. 我们现在可以发送一张测试图像,并观察推理脚本如何处理 Docker 日志中的推理请求:

    import requests
    import json 
    payload = trainset[0][0].numpy().tobytes()
    url = 'http://127.0.0.1:8080/invocations'
    content_type = 'application/x-npy'
    accept_type = "application/json"
    headers = {'content-type': content_type, 'accept': accept_type}
    response = requests.post(url, data=payload, headers=headers)
    print(json.loads(response.content)[0])
    

在前面的代码块中,我们执行了以下操作:

  • 构造推理有效载荷并将其序列化为bytes对象

  • 构造了content-typeaccept-type HTTP 头,指示推理服务器客户端发送的内容类型以及期望的内容类型

  • 向本地 SageMaker 端点发送请求

  • 读取响应输出

如果出现任何问题,您可以登录到运行中的推理容器中,检查运行时环境,或使用您的 IDE 功能设置调试会话。

总结

在本章中,我们回顾了一些可用的解决方案和最佳实践,讲解了如何为 Amazon SageMaker 组织深度学习代码的开发。根据您的使用案例需求和个人偏好,您可以选择在本地创建 DIY 环境,或使用 SageMaker 的笔记本环境之一——笔记本实例或 Studio 笔记本。您还学习了如何在本地测试 SageMaker 深度学习容器,以加速开发过程并避免额外的测试费用。

在下一章中,我们将重点介绍 SageMaker 的数据管理和数据处理。由于许多深度学习(DL)问题的训练数据集较大,并且需要进行预处理或后处理,因此理解最佳的存储解决方案至关重要。我们还将讨论使用 SageMaker 功能进行数据标注和数据处理的各个方面,以及访问训练数据的最佳实践。

第四章:管理深度学习数据集

深度学习模型通常需要大量的训练数据来学习有用的模式。在许多现实应用中,新数据会持续收集、处理,并添加到训练数据集中,以便您的模型可以定期进行重新训练,从而适应不断变化的现实环境。在本章中,我们将探讨 SageMaker 的能力以及其他 AWS 服务,帮助您管理训练数据。

SageMaker 提供了广泛的集成功能,您可以使用 AWS 通用数据存储服务,如 Amazon S3、Amazon EFS 和 Amazon FSx for Lustre。此外,SageMaker 还有专为机器学习ML)设计的存储解决方案——SageMaker Feature Store。我们将讨论根据数据类型、消费和摄取模式选择存储解决方案的时机。

在许多情况下,在使用训练数据之前,您需要对其进行预处理。例如,数据需要转换成特定格式,或数据集需要通过修改后的样本版本进行增强。本章将回顾 SageMaker Processing,以及如何使用它来处理大规模的机器学习(ML)数据集。

本章最后,我们将探讨一些先进的技术,如何利用 AWS 数据流工具优化 TensorFlow 和 PyTorch 模型的数据检索过程。

本章将涵盖以下主题:

  • 为 ML 数据集选择存储解决方案

  • 大规模处理数据

  • 优化数据存储和检索

阅读完本章后,您将知道如何组织您的深度学习(DL)数据集的生命周期,以便在 SageMaker 上进行训练和推理。我们还将通过一些动手示例来帮助您获得数据处理和数据检索方面的实际技能。

技术要求

在本章中,我们将提供代码示例,以便您可以开发实际技能。完整的代码示例可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter4/查看。

要跟随这段代码,您需要以下内容:

  • AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户

  • 已建立的 SageMaker 笔记本、SageMaker Studio 笔记本或本地兼容 SageMaker 环境

为 ML 数据集选择存储解决方案

AWS Cloud 提供了一系列可以用来存储推理和训练数据的存储解决方案。在选择最佳存储解决方案时,您可以考虑以下因素:

  • 数据量和速度

  • 数据类型及相关元数据

  • 消费模式

  • 备份和保留要求

  • 安全性和审计要求

  • 集成能力

  • 存储、写入和读取数据的价格

仔细分析你的具体需求可能会为你的用例建议正确的解决方案。通常,在数据生命周期的不同阶段会组合使用几种存储解决方案。例如,你可以将用于推理的低延迟需求数据存储在更快但更昂贵的存储中;然后,将数据移动到更便宜、更慢的存储解决方案中用于训练和长期保存。

存在几种常见的存储类型,它们具有不同的特性:文件系统、对象存储和块存储解决方案。Amazon 为每种存储类型提供托管服务。我们将在以下小节中回顾它们的特性以及如何在 SageMaker 工作负载中使用它们。接下来,我们将重点介绍 Amazon SageMaker 特征库,因为它为 ML 工作负载和数据集提供了几个独特的功能。

Amazon EBS – 高性能块存储

块存储解决方案旨在快速检索和操作数据。数据在物理设备上被分割成块,以便高效利用。块存储使你能够将数据从运行时环境中抽象和解耦。在数据检索时,存储解决方案将这些块重新组合并返回给用户。

Amazon EBS 是一种完全托管的块存储解决方案,支持多种用例,适用于不同的读写模式、吞吐量和延迟要求。Amazon EBS 的一个主要用例是作为附加到 Amazon EC2 计算节点的数据卷。

Amazon SageMaker 提供与 EBS 的无缝集成。在以下示例中,我们正在为四个节点配置一个训练任务;每个节点将附加一个 100 GB 的 EBS 卷。训练数据将由 SageMaker 下载并存储在 EBS 卷上:

from sagemaker.pytorch import PyTorch
pytorch_estimator = PyTorch(
                        session=session,
                        entry_point=f'path/to/train.py',
                        role=role,
                        instance_type="ml.m4.xlarge",
  volume_size=100,
                        instance_count=4,
                        framework_version="1.9.0",
                        )

请注意,你无法自定义使用的 EBS 卷类型。仅支持通用 SSD 卷。一旦训练任务完成,所有实例和附加的 EBS 卷将被清除。

Amazon S3 – 行业标准对象存储

对象存储实现了一个平面结构,每个文件对象都有一个唯一标识符(以路径形式表示)和关联的数据对象。平面结构允许你线性扩展对象存储解决方案,并在保持低成本的同时,支撑高吞吐量的数据读写。

对象存储可以处理不同类型和大小的对象。对象存储还允许你存储与每个对象相关的元数据。数据读写通常通过 HTTP API 进行,这使得集成变得更加容易。但需要注意的是,对象存储解决方案通常比文件系统或块存储解决方案慢。

亚马逊 S3 是首个 PB 级别的云对象存储服务。它提供持久性、可用性、性能、安全性和几乎无限的可扩展性,同时成本非常低廉。许多对象存储解决方案都遵循亚马逊 S3 的 API。亚马逊 S3 用于存储客户数据,但它也用于许多 AWS 内部功能和服务,其中数据需要持久化。

SageMaker 提供了与 Amazon S3 的无缝集成,用于存储输入和输出对象,如数据集、日志流、作业输出和模型工件。让我们看一个训练作业示例,以及如何定义我们将存储输入和输出的位置:

  • model_uri 参数指定了模型工件的 S3 位置,例如预训练权重、分词器等。SageMaker 会自动将这些工件下载到每个训练节点。

  • checkpoint_s3_uri 定义了训练过程中上传训练检查点的 S3 位置。请注意,实现检查点功能由开发者在训练脚本中负责。

  • output_path 指定了 SageMaker 在训练作业完成后,从训练节点上传的所有输出工件的 S3 目标位置。

  • tensorboard_output_config 定义了在 S3 上存储 TensorBoard 日志的位置。请注意,SageMaker 会在训练作业执行期间持续上传这些日志,因此你可以在 TensorBoard 中实时监控训练进度:

    from sagemaker.huggingface.estimator import HuggingFace
    from sagemaker.debugger import TensorBoardOutputConfig
    from sagemaker import get_execution_role
    role=get_execution_role()
    estimator = HuggingFace(
        py_version="py36",
        entry_point="train.py",
        pytorch_version="1.7.1",
        transformers_version="4.6.1",
        instance_type="ml.p2.xlarge",
        instance_count=2,
        role=role,
        model_uri="s3://unique/path/models/pretrained-bert/",
        checkpoint_s3_uri="s3://unique/path/training/checkpoints",
        output_path="s3://unique/path/training/output",
        tensorboard_output_config = TensorBoardOutputConfig(
            s3_output_path='s3://unique/path/tensorboard/ ',
            container_local_output_path='/local/path/'
            ),
    )
    

我们还使用 S3 来存储我们的训练数据集。请看以下 estimator.fit() 方法,它定义了我们训练数据的位置:

estimator.fit({
    "train":"s3://unique/path/train_files/",
    "test":"s3://unique/path/test_files"}
    )

在这里,"train""test" 参数被称为 /opm/ml/input/data/{channel_name} 目录。此外,训练工具包将创建 SM_CHANNEL_{channel_name} 环境变量,您可以在训练脚本中使用这些变量来本地访问模型工件。

如上面代码块所示,Amazon S3 可用于存储 SageMaker 训练作业的输入和输出工件。

文件、快速文件和管道模式

S3 存储是一个常用的训练数据集存储位置。默认情况下,当使用存储在 S3 上的数据时,所有与路径匹配的对象将被下载到每个计算节点,并存储在其 EBS 卷中。这被称为 文件模式。

然而,在许多场景中,训练数据集可能有数百 GB 或更大。即使在训练开始之前,下载如此大的文件也需要相当长的时间。为了减少启动训练作业所需的时间,SageMaker 支持 管道模式,允许你从 S3 位置流式传输数据,而无需完全下载。这使你能够立即启动训练作业,并在训练过程中根据需要获取数据批次。

Pipe 模式的一个缺点是需要使用框架特定的实现来流式传输数据。最近引入的FastFile模式解决了这一问题。FastFile 模式允许您直接从 S3 流式传输数据,而无需实现任何特定的数据加载器。在您的训练或处理脚本中,您可以将FastFile视为存储在磁盘上的常规文件;Amazon SageMaker 会为您处理读写操作。

我们将在优化数据存储与检索部分中,学习如何使用FastFilePipe模式组织 S3 流式传输的训练代码。

FullyReplicated 和 ShardedByKey

在许多训练和数据处理任务中,我们希望将作业并行化到多个计算节点上。在有大量数据对象的场景中,我们可以通过将完整的对象集拆分成唯一的子集来拆分任务。

为了实现这种场景,SageMaker 支持ShardedByKey模式,该模式尝试均匀地拆分所有匹配对象,并将每个节点分配一个独特的对象子集。例如,如果您的数据集中有 n 个对象,而作业中有 k 个计算节点,那么每个计算节点将获得一个独特的 n/k 个对象的子集。

除非另有说明,默认模式FullyReplicated会在 SageMaker 下载所有匹配对象到所有节点时使用。

我们将在分布式数据处理部分中学习如何分配数据处理任务的实际技能。

Amazon EFS – 通用共享文件系统

Amazon EFS 是一种托管的文件存储服务,易于设置并自动扩展到 PB 级别。它提供文件系统接口和文件语义,如文件锁定和强一致性。与 Amazon EBS 不同,EBS 允许您将存储附加到单个计算节点,而 EFS 可以同时被数百或数千个计算节点使用。这使您能够组织高效的数据共享,而无需重复和分发数据。

Amazon SageMaker 允许您使用 EFS 存储训练数据集。以下代码展示了如何在训练作业配置中使用 Amazon EFS,通过 FileSystemInput 类进行配置。请注意,在此情况下,我们已配置为数据的只读访问(file_system_access_mode 参数的 ro 标志),这通常是训练作业的情况。不过,您也可以通过将 file_system_access_mode 设置为 rw 来指定读写权限。

from sagemaker.tensorflow.estimator import TensorFlow
from sagemaker.inputs import FileSystemInput
from sagemaker import get_execution_role
role=get_execution_role()
estimator = TensorFlow(entry_point='train.py',
                       role=role,
                       image_uri="image/uri",
                       instance_count=4,
                       instance_type='ml.c4.xlarge')
file_system_input = FileSystemInput(file_system_id='fs-1',
                                    file_system_type='EFS',
                                    directory_path='/tensorflow',
                                    file_system_access_mode='ro')
estimator.fit(file_system_input)

在这里,您可以控制其他 EFS 资源。根据数据读写的延迟要求,您可以从几种模式中选择,这些模式定义了文件系统的延迟和并发特性。在本书编写时,EFS 可以支持每秒 10+ GB 的吞吐量,并可扩展到数千个连接的计算节点。

Amazon FSx for Lustre – 高性能文件系统

Amazon FSx for Lustre 是一种针对机器学习和高性能计算HPC)工作负载优化的文件存储服务。它设计用于子毫秒延迟的读写操作,并且能够提供数百 GB/s 的吞吐量。你还可以选择将数据存储在 S3 中,并与 Amazon FSx for Lustre 文件系统进行同步。在这种情况下,FSx 系统会将 S3 对象呈现为文件,并允许你将数据更新回 S3 原点。

亚马逊 SageMaker 支持将训练数据存储在 FSx for Lustre 文件系统中。训练作业配置与使用 EFS 文件系统类似;唯一的区别是 file_system_type 参数设置为 FSxLustre。以下代码展示了一个示例训练作业:

from sagemaker.inputs import FileSystemInput
from sagemaker import get_execution_role
role=get_execution_role()
estimator = TensorFlow(entry_point='train.py',
                       role=role,
                       image_uri="image/uri",
                       instance_count=4,
                       instance_type='ml.c4.xlarge')
file_system_input = FileSystemInput(
                       file_system_id='fs-XYZ', 
                       file_system_type='FSxLustre',
                       directory_path='/tensorflow',
                       file_system_access_mode='ro')
estimator.fit(file_system_input)

请注意,在配置你的 Lustre 文件系统时,你可以选择 SSD 或 HDD 存储。对于对延迟敏感的工作负载,应该选择 SSD;而对于具有高吞吐量需求的工作负载,HDD 更为适合。

SageMaker 特征存储——专为机器学习设计的存储

到目前为止,我们讨论了可以用于存储数据的通用文件和对象存储服务,这些服务可以应用于你的 SageMaker 工作负载。然而,现实中的机器学习工作流在特征工程和数据管理方面可能会遇到一些挑战,如下所示:

  • 管理数据摄取管道,以保持数据的最新状态

  • 在组织内不同团队之间组织数据使用,并消除重复工作

  • 在推理和训练工作负载之间共享数据(在需要时)

  • 管理数据集一致性、元数据和版本控制

  • 数据的临时分析

为了解决这些挑战,SageMaker 提供了一种专门为机器学习设计的数据存储解决方案,称为特征存储。它通过减少重复步骤并提供一组 API 来获取、转换和使用数据进行推理和模型训练,从而加速数据处理和整理。

它的核心概念是特征——数据记录的单一属性。每个数据记录由一个或多个特征组成。此外,数据记录还包含元数据,如记录更新时间、唯一记录 ID 和状态(是否已删除)。特征可以是字符串、整数或小数类型。数据记录及其相关特征可以组织成逻辑单元,称为特征组

现在,让我们回顾一下 SageMaker 特征存储的关键功能。

在线存储和离线存储

特征存储支持多种存储选项,以适应不同的使用场景:

  • 离线存储旨在存储数据,其场景中数据检索延迟不关键,例如存储训练或批量推理的数据。你的数据集存储在 S3 中,可以使用 Amazon Athena SQL 引擎进行查询。

  • 在线存储允许你以毫秒级延迟检索单个或一批记录,适用于实时推理场景。

  • 离线和在线存储允许你将相同的数据存储在两种形式的存储中,并在推理和训练场景中使用。

导入接口

有几种方法可以将数据导入 Feature Store。其中一种方法是使用 Feature Store 的 PutRecord API,它允许您写入单条或批量记录。这将把记录写入离线存储和在线存储。

另一种选择是使用 Spark 连接器。如果您已经有基于 Spark 的数据处理管道,这是一种方便的数据导入方式。

分析查询

当数据存储在离线存储中时,可以使用 Athena SQL 通过 SQL 语法查询数据集。这对于拥有不同编码技能水平的多元化团队非常有帮助。由于 Feature Store 包含有用的元数据字段,如事件时间状态,您可以使用这些时间进行时间旅行查询,例如,获取数据集在某一时刻的历史快照。

特征发现

一旦数据被导入 Feature Store,您可以通过 SageMaker Studio 使用直观的 UI 组件查看和分析数据集,而无需编写任何代码:

图 4.1 – 通过 SageMaker Studio UI 发现 Feature Store 数据集

图 4.1 – 通过 SageMaker Studio UI 发现 Feature Store 数据集

现在我们了解了 Feature Store 相比于更通用的存储解决方案的价值主张,接下来让我们看看它如何在典型的深度学习场景中使用,当我们希望将分词文本与原始文本并排存放时。

使用 Feature Store 进行推理和训练

在这个实际示例中,我们将学习如何使用 SageMaker Feature Store 导入、处理和消费包含 IMDb 评论的数据集。我们将采用包含评论的原始数据集,并运行自定义 BERT 分词器将非结构化文本转换为一组整数令牌。然后,我们将把包含分词文本特征的数据集导入 Feature Store,这样下次使用时就无需再进行分词了。之后,我们将训练模型以分类正面和负面评论。

我们将使用 SageMaker Feature Store SDK 与 Feature Store APIs 进行交互。我们将使用 HuggingFace Datasets(huggingface.co/docs/datasets/)和 Transformers(huggingface.co/docs/transformers/index)库来对文本进行分词,并进行训练和推理。请确保已安装这些库。

准备数据

按照以下步骤准备数据:

  1. 第一步是获取包含 IMDb 评论的初始数据集:

    from datasets import load_dataset
    dataset = load_dataset("imdb")
    
  2. 然后,我们必须将数据集转换为与EventTimeID兼容的 pandas DataFrame。这两个字段是 Feature Store 所必需的,以支持快速检索和特征版本控制:

    import pandas as pd
    import time
    dataset_df = dataset['train'].to_pandas()
    current_time_sec = int(round(time.time()))
    dataset_df["EventTime"] = pd.Series([current_time_sec]*len(dataset_df), dtype="float64")
    dataset_df["ID"] = dataset_df.index
    dataset_df["text"] = dataset_df["text"].astype('string')
    dataset_df["text"] = dataset_df["text"].str.encode("utf8")
    dataset_df["text"] = dataset_df["text"].astype('string')
    
  3. 现在,让我们运行下载的 Distilbert 模型预训练的分词器,并向数据集添加一个新属性 tokenized-text。请注意,我们将 tokenized-text 转换为字符串,因为 SageMaker Feature Store 不支持集合数据类型,如数组或映射:

    from transformers import DistilBertTokenizerFast
    tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
    dataset_df["tokenized-text"] = tokenizer(dataset_df["text"].tolist(), truncation=True, padding=True)["input_ids"]
    dataset_df["tokenized-text"] = dataset_df["tokenized-text"].astype('string')
    

结果,我们得到了一个 pandas DataFrame 对象,里面包含了我们希望摄取到 Feature Store 中的特征。

数据摄取

下一步是为 Feature Store 配置资源并准备摄取。请按照以下步骤操作:

  1. 我们将从配置特征组并准备特征定义开始。请注意,由于我们将数据集存储在 pandas DataFrame 中,Feature Store 可以使用该 DataFrame 推断特征类型:

    from sagemaker.feature_store.feature_group import FeatureGroup
    imdb_feature_group_name = "imdb-reviews-tokenized"
    imdb_feature_group = FeatureGroup(name=imdb_feature_group_name, sagemaker_session=sagemaker_session)
    imdb_feature_group.load_feature_definitions(data_frame=dataset_df)
    
  2. 现在我们已经准备好了特征组配置,接下来可以创建它。这个过程可能需要几分钟,所以我们来加一个 Waiter。由于我们打算同时使用在线存储和离线存储,我们将 enable_online_store 标志设置为 True

    imdb_feature_group.create(    
    s3_uri=f"s3://{s3_bucket_name}/{imdb_feature_group_name}",
        record_identifier_name="ID",
        event_time_feature_name="EventTime",
        role_arn=role,
        enable_online_store=True
    )
    # Waiter for FeatureGroup creation
    def wait_for_feature_group_creation_complete(feature_group):
        status = feature_group.describe().get('FeatureGroupStatus')
        print(f'Initial status: {status}')
        while status == 'Creating':
            print(f'Waiting for feature group: {feature_group.name} to be created ...')
            time.sleep(5)
            status = feature_group.describe().get('FeatureGroupStatus')
        if status != 'Created':
            raise SystemExit(f'Failed to create feature group {feature_group.name}: {status}')
        print(f'FeatureGroup {feature_group.name} was successfully created.')
    wait_for_feature_group_creation_complete(imdb_feature_group)
    
  3. 一旦组可用,我们就可以准备摄取数据。由于我们有完整的数据集,我们将使用批量摄取 API,如下所示:

    imdb_feature_group.ingest(data_frame=dataset_df, max_processes=16, wait=True)
    
  4. 数据被摄取后,我们可以运行一些分析查询。例如,我们可以检查数据集是平衡的还是不平衡的。如前所述,Feature Store 支持使用 Amazon Athena SQL 引擎查询数据:

    athena_query = imdb_feature_group.athena_query()
    imdb_table_name = athena_query.table_name
    result = athena_query.run(f'SELECT "label", COUNT("label") as "Count" FROM "sagemaker_featurestore"."{imdb_table_name}" group by "label";', output_location=f"s3://{s3_bucket_name}/athena_output")
    athena_query.wait()
    print(f"Counting labels in dataset: \n {athena_query.as_dataframe()}")
    

运行这个过程需要一些时间,但最终你应该能得到我们数据集中标签的数量。

使用 Feature Store 进行训练

现在我们有了可用的数据,接下来让我们训练二分类模型。由于 Feature Store 中的数据以 Parquet 格式存储在指定的 S3 位置(parquet.apache.org/),我们可以直接使用 Parquet 文件进行训练。

为了处理 Parquet 文件,我们需要确保数据读取器能够识别这种格式。为此,我们可以使用 pandas 的 .read_parquet() 方法。然后,我们可以将 pandas DataFrame 对象转换为 HuggingFace 数据集,并选择我们关心的属性——tokenized-textlabel

    df = pd.read_parquet(args.training_dir)
    df["input_ids"] = df["tokenized-text"].astype("string")
    train_dataset = Dataset.from_pandas(df[["input_ids", "label"]])

现在,我们需要将 tokenized-text 从字符串转换为整数列表:

    def string_to_list(example):
        list_of_str = example["input_ids"].strip("][").split(", ")
        example["input_ids"] = [int(el) for el in list_of_str]
        return example
    train_dataset = train_dataset.map(string_to_list)

其余的训练脚本是相同的。你可以在 github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter4/1_sources/train.py 找到完整代码。

现在我们已经修改了训练脚本,准备好运行训练任务了:

  1. 首先,我们必须获取数据集的位置:

    train_dataset_uri = imdb_feature_group.describe()['OfflineStoreConfig']["S3StorageConfig"]["ResolvedOutputS3Uri"]
    
  2. 现在,我们必须将它传递给我们的 Estimator 对象:

    from sagemaker.huggingface.estimator import HuggingFace
    estimator = HuggingFace(
        py_version="py36",
        entry_point="train.py",
        source_dir="1_sources",
        pytorch_version="1.7.1",
        transformers_version="4.6.1",
        hyperparameters={
            "model_name":"distilbert-base-uncased",
            "train_batch_size": 16,
            "epochs": 3
            # "max_steps": 100 # to shorten training cycle, remove in real scenario
        },
        instance_type="ml.p2.xlarge",
        debugger_hook_config=False,
        disable_profiler=True,
        instance_count=1,
        role=role
    )
    estimator.fit(train_dataset_uri)
    

经过一段时间(取决于你使用的 epoch 或步骤数),模型应该已经训练好,可以根据输入的文本对评论进行分类。

使用 Feature Store 进行推理

对于推理,我们可以使用 Boto3 库中的 Feature Store 运行时客户端来获取单个记录或批量记录:

import boto3
client = boto3.client('sagemaker-featurestore-runtime')

请注意,你需要知道记录的唯一 ID 才能检索它们:

response = client.batch_get_record(
    Identifiers=[
        {
            'FeatureGroupName':imdb_feature_group.name,
            'RecordIdentifiersValueAsString': ["0", "1", "2"], # picking several records to run inference.
            'FeatureNames': [
                'tokenized-text', "label", 'text'
            ]
        },
    ]
)
# preparing the inference payload
labels = []
input_ids = []
texts = []
for record in response["Records"]:
    for feature in record["Record"]:
        if feature["FeatureName"]=="label":
            labels.append(feature["ValueAsString"])
        if feature["FeatureName"]=="tokenized-text":
            list_of_str = feature["ValueAsString"].strip("][").split(", ")
            input_ids.append([int(el) for el in list_of_str])
        if feature["FeatureName"]=="text":
            texts.append(feature["ValueAsString"])    

现在,你可以将这个推理请求发送到你已部署的模型。请参考以下笔记本,查看完整的端到端示例:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter4/1_Managing_data_in_FeatureStore.ipynb

在本节中,我们回顾了可以用于存储 ML 数据以供推理和训练的选项。但数据“原样”使用的情况很少见。在许多场景中,在将数据用于 ML 工作负载之前,你需要在大规模上持续处理数据。SageMaker Processing 提供了一个可扩展且灵活的机制,用于大规模处理数据。我们来看看。

大规模处理数据

SageMaker Processing 允许你在云中运行容器化代码。这对于数据的预处理和后处理、特征工程以及模型评估等场景非常有用。SageMaker Processing 对于临时工作负载以及周期性任务都很有用。

与训练任务一样,Amazon SageMaker 提供了一个托管的底层计算和数据基础设施体验。你需要提供一个处理任务配置、代码和要使用的容器,但 SageMaker 将负责配置实例、部署容器化代码,并运行和监控任务及其进度。一旦任务达到终止状态(成功或失败),SageMaker 将把结果工件上传到 S3 存储并撤销集群。

SageMaker Processing 提供了两个预构建容器:

  • 一个具有运行 Spark 计算依赖关系的 PySpark 容器

  • 一个 scikit-learn 容器

在选择内置处理容器时,请注意,PySpark 容器支持分布式 Spark 任务。它允许你在 Spark 集群中协调分布式数据处理,维护数据集中的全局状态,并通过 Spark UI 可视化处理任务。同时,scikit-learn 容器不支持共享全局状态,因此每个处理节点独立运行。可以通过将数据集分割成子数据集并独立处理每个子数据集,进行有限的任务协调。

你还可以提供一个自带容器BYO)处理容器,几乎可以使用任何运行时配置来运行 SageMaker Processing。这个灵活性使你能够轻松地将现有的处理代码迁移到 SageMaker Processing 上,几乎无需任何额外努力:

图 4.2 – SageMaker Processing 节点

图 4.2 – SageMaker Processing 节点

让我们尝试构建一个处理容器,并运行一个多节点处理任务,以增强图像数据集以便进一步训练。

使用 SageMaker Processing 增强图像数据

在这个示例中,我们将从 Kaggle 下载 325 种鸟类数据集(www.kaggle.com/gpiosenka/100-bird-species/)。然后,我们将通过修改图像(旋转、裁剪、调整大小)来增强这个数据集,以提高后续图像分类任务的性能。为了进行图像转换,我们将使用 Keras 库。接下来,我们将在多个节点上运行处理作业,以加速任务的执行。请按照以下步骤操作:

  1. 我们将从构建自定义处理容器开始。请注意,SageMaker 使用 docker run image_uri 命令运行处理容器,因此我们需要在 Dockerfile 中指定入口点。我们使用官方的 Python 3.7 容器,并配备基本的 Debian 版本:

    FROM python:3.7-slim-buster
    ########### Installing packages ##########
    RUN pip3 install pandas numpy tensorflow numpy scipy
    RUN pip install Pillow
    ENV PYTHONUNBUFFERED=TRUE
    ########### Configure processing scripts ##########
    ARG code_dir=/opt/ml/code
    RUN mkdir -p $code_dir
    COPY 2_sources $code_dir
    WORKDIR $code_dir
    ENTRYPOINT ["python3","processing.py"]
    

我们将从构建自定义处理容器开始。

  1. 现在,我们需要提供我们的处理代码。我们将使用 keras.utils 将原始数据集加载到内存中,并指定必要的转换:

        dataset = keras.utils.image_dataset_from_directory(
            args.data_location,
            labels="inferred",
            label_mode="int",
            class_names=None,
            color_mode="rgb",
            batch_size=args.batch_size,
            image_size=(WIDTH, HEIGHT),
            shuffle=True,
            seed=None,
            validation_split=None,
            subset=None,
            interpolation="bilinear",
            follow_links=False,
            crop_to_aspect_ratio=False,
        )
        datagen = ImageDataGenerator(
            rotation_range=40,
            width_shift_range=0.2,
            height_shift_range=0.2,
            shear_range=0.2,
            zoom_range=0.2,
            horizontal_flip=True,
            fill_mode="nearest",
        )
    
  2. 由于 Keras 生成器在内存中操作,我们需要将生成的图像保存到磁盘:

        for batch_data, batch_labels in dataset.as_numpy_iterator():
            print(f"Processing batch with index {i} out from {len(dataset)}")
            for image, label in zip(batch_data, batch_labels):
                label_name = class_lookup.iloc[label]["class"]
                image_save_dir = os.path.join(augmented_root_dir, label_name)
                os.makedirs(image_save_dir, exist_ok=True)
                j = 0
                image = np.expand_dims(image, axis=0)
                # generate 5 new augmented images
                for batch in datagen.flow(
                    image,
                    batch_size=1,
                    save_to_dir=image_save_dir,
                    save_prefix="augmented",
                    save_format="jpeg",
                ):
                    j += 1
                    if j > max_augmentations:
                        break
            i += 1
            if args.max_samples is not None:
                if i > args.max_samples:
                    break
    

我们将增强后的图像保存在类似的目录结构中,其中标签由目录名称定义。

  1. 一旦我们拥有了 BYO 容器和处理代码,就可以准备调度处理作业。首先,我们需要用基本的作业配置实例化 Processor 对象,例如实例的数量和类型以及容器镜像:

    from sagemaker.processing import Processor, ProcessingInput, ProcessingOutput
    lookup_location = "/opt/ml/processing/lookup"
    data_location = "/opt/ml/processing/input"
    output_location = '/opt/ml/processing/output'
    sklearn_processor = Processor(
                          image_uri=image_uri,
                          role=role,
                          instance_count=2,
                          base_job_name="augmentation",
                          sagemaker_session=sess, 
                          instance_type="ml.m5.xlarge")
    

要启动作业,我们必须执行 .run() 方法。此方法允许我们提供额外的配置参数。例如,为了均匀分配任务,我们需要将数据集拆分为多个块。使用 ShardedByKey 分发类型可以轻松实现这一点。在这种情况下,SageMaker 将尝试在我们的处理节点之间均匀分配对象。SageMaker Processing 允许你通过 arguments 集合传递自定义脚本配置。你需要确保处理脚本能够正确解析这些命令行参数:

sklearn_processor.run(
    inputs=[
      ProcessingInput(
          source=dataset_uri,
          destination=data_location,
          s3_data_distribution_type="ShardedByS3Key"),
      ProcessingInput(
          source=class_dict_uri,
          destination=lookup_location),],
    outputs=[
      ProcessingOutput(
          source=output_location)],
          arguments = [
               "--data_location", data_location, 
               "--lookup_location", lookup_location,
               "--output_location", output_location,
               "--batch_size", "32",
               "--max_samples", "10",
               "--max_augmentations", "5"
               ])                     

有关完整的处理代码,请参考github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter4/2_sources/processing.py

这个示例应能帮助你直观了解 SageMaker Processing 如何用于数据处理需求。同时,SageMaker Processing 足够灵活,可以运行任何任意任务,比如批量推断、数据聚合与分析等。

在下一部分,我们将讨论如何优化大型深度学习数据集的数据存储和检索。

优化数据存储和检索

在训练SOTA DL模型时,通常需要一个大型数据集来进行训练。存储和检索这样的大型数据集可能会很昂贵。例如,流行的计算机视觉数据集COCO2017大约有 30 GB,而用于 NLP 任务的Common Crawl数据集则有数百 TB。处理如此庞大的数据集需要仔细考虑存储数据集的位置以及如何在推理或训练时检索它。在本节中,我们将讨论一些优化策略,帮助您在选择存储和检索策略时做出决策。

选择存储解决方案

在选择最优存储方案时,您可以考虑以下因素,除此之外还有其他因素:

  • 存储和数据检索的费用

  • 数据检索的延迟和吞吐量要求

  • 数据分区

  • 数据刷新频率

让我们来看一下各种存储解决方案的优缺点:

  • Amazon S3 提供了在考虑的存储方案中最便宜的选择。然而,您应该意识到,Amazon S3 还会对数据传输和数据请求收取费用。在您的数据集包含大量小文件的情况下,您可能会因为PUTGET操作记录而产生相当可观的费用。您可以考虑将小对象批量合并成大对象以减少这部分费用。请注意,从另一个 AWS 区域检索数据时会产生额外的费用。将工作负载和数据放在同一 AWS 区域可能是避免这些费用的合理选择。S3 通常也是最慢的存储解决方案。默认情况下,Amazon SageMaker 会在训练开始之前从 S3 下载所有对象。这个初始下载时间可能需要几分钟,并会增加整体训练时间。例如,在COCO2017数据集的情况下,从 S3 下载到训练节点大约需要 20 分钟。亚马逊提供了几种机制,可以直接从 S3 流式传输数据,避免下载时间。我们将在本节中讨论这些机制。

  • Amazon EFS存储通常比 Amazon S3 更昂贵。然而,与 Amazon S3 不同,Amazon EFS 不会对读写操作收取任何费用。由于 EFS 提供了文件系统接口,计算节点可以直接挂载到包含数据集的 EFS 目录,并立即使用它,无需下载数据集。Amazon EFS 提供了一个便捷的机制,让不同的工作负载或团队之间共享可重用的数据集。

  • Amazon FSx for Lustre 提供了最低的延迟,但也具有最高的存储成本。像 Amazon EFS 一样,它不需要任何下载时间。常见的使用场景之一是将数据存储在 S3 中。当你需要运行实验集时,可以通过从 S3 同步来配置 FSx for Lustre,这样可以无缝地将数据从 S3 复制到你的文件系统中。之后,你可以运行实验,并将 FSx for Lustre 作为数据源,利用最低的延迟进行数据检索。实验完成后,可以解除配置 Lustre 文件系统,以避免任何额外的成本,同时将原始数据保留在 S3 中。

  • SageMaker 特征存储具有最全面的机器学习特定功能;但是,它也有其缺点和强假设。由于其离线存储由 S3 支持,因此它具有类似的成本结构和延迟考虑。在线存储会增加额外的存储、读取和写入成本。SageMaker 特征存储适合在需要重用相同数据集进行推理和训练工作负载的场景中。特征存储的另一个常见用例是当你需要进行审计要求或对数据集进行分析查询时。请注意,由于特征存储仅支持有限的数据类型(例如,它不支持任何集合类型),因此在从特征存储消费数据时,可能需要进行类型转换。

AWS 提供了广泛的存储解决方案,有时选择哪种解决方案可能不太明显。像往常一样,重要的是从理解你的使用案例需求和成功标准开始(例如,最低的延迟、最高的吞吐量或最具成本效益的解决方案)。

流式数据集

Amazon S3 是一个流行的大型机器学习数据集存储解决方案,因其低成本、高耐用性、便捷的 API 以及与其他服务(如 SageMaker)的集成。在前面讨论的部分中,我们提到使用 S3 存储训练数据集的一个缺点是,训练开始之前需要将数据集下载到训练节点。

你可以选择使用 ShardedByKey 分发策略,这将减少每个训练节点下载的数据量。然而,这种方法仅减少了需要下载到训练节点的数据量。对于大型数据集(100 GB 以上),它仅部分解决了问题。你还需要确保训练节点有足够的 EBS 卷容量来存储数据。

减少训练时间的另一种方法是直接从 Amazon S3 流式传输数据,而不是事先下载数据。Amazon SageMaker 提供了几种 S3 数据流式传输的实现方式:

  • 针对特定框架的流式传输实现,例如 TensorFlow 的 PipeModeDataset 和 PyTorch 的 Amazon S3 插件

  • 框架无关的 FastFile 模式

让我们回顾一下这些方法的优点。

TensorFlow 的 PipeModeDataset

使用PipeModeDataset时,你的训练程序可以直接从 S3 读取数据,而无需管理对 S3 对象的访问。在使用PipeModeDataset时,需确保你使用的是与之匹配的 TensorFlow 版本。

在配置 SageMaker 训练作业时启用 SageMaker Pipe 模式。如果你将多个数据集存储在同一个 S3 路径下,可以将它们映射到单个管道。请注意,SageMaker 最多支持 20 个管道。如果需要超过 20 个管道,你可以考虑使用增强的清单文件,它允许你显式列出一组 S3 对象进行流式传输。在训练过程中,SageMaker 将从清单文件中读取对象并将它们流式传输到管道中。

PipeModeDataset支持以下数据集格式:文本行、RecordIO 和 TFRecord。如果你有其他格式的数据集(例如单独的图像文件),则需要将数据集转换为支持的格式。请注意,PipeModeDataset的性能受到文件数量和大小的影响。通常建议将文件大小保持在 100 到 200 MB 之间,以获得最佳性能。

注意

由于PipeModeDataset实现了 TensorFlow Dataset API,你可以使用熟悉的方法来操作数据集,例如.apply().map()PipeModeDataset也可以直接传递给 TensorFlow Estimator。

PipeModeDataset与 TensorFlow Dataset 之间有几个不同之处,你需要考虑以下几点:

  • PipeModeDataset按顺序从文件中读取数据。SageMaker 支持ShuffleConfigdocs.aws.amazon.com/sagemaker/latest/APIReference/API_ShuffleConfig.xhtml)参数,用于打乱读取文件的顺序。你还可以调用.shuffle()方法进一步打乱记录顺序。

  • PipeModeDataset仅支持三种数据类型,所有数据都需要转换为支持的格式之一。

  • PipeModeDataset在训练时操作数据的控制功能有限。例如,如果你需要增强分类数据集中低频类别的样本,你将需要使用单独的管道来流式传输低频类别的样本,并在训练脚本中处理增强过程。

  • PipeModeDataset不支持 SageMaker 本地模式,因此调试训练程序可能会比较困难。在使用 SageMaker Pipe 模式时,你无法访问 SageMaker 如何将数据对象流式传输到管道中的内部实现。

让我们看一下如何使用PipeModeDataset。在这个例子中,为了训练目的,我们将 CIFAR-100 数据集转换为 TFRecords,然后在训练时通过PipeModeDataset流式传输此数据集。为了简洁起见,我们将提供一个编辑版,而不是列出整个例子。完整源代码可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter4/3_Streaming_S3_Data.ipynb查看。请按照以下步骤操作:

  1. 让我们首先将数据集转换为 TFRecord 格式。在以下代码块中,包含一个方法,该方法遍历一批文件,将一对图像和标签转换为 TensorFlow 的Example类,并将一批Example对象写入一个单独的TFRecord文件中:

    def convert_to_tfrecord(input_files, output_file):
        """Converts a file to TFRecords."""
        print("Generating %s" % output_file)
        with tf.io.TFRecordWriter(output_file) as record_writer:
            for input_file in input_files:
                data_dict = read_pickle_from_file(input_file)
                data = data_dict[b"data"]
                labels = data_dict[b"fine_labels"]
                num_entries_in_batch = len(labels)
                for i in range(num_entries_in_batch):
                    example = tf.train.Example(
                        features=tf.train.Features(
                            feature={
                                "image": _bytes_feature(data[i].tobytes()),
                                "label": _int64_feature(labels[i]),
                            }
                        )
                    )
                    record_writer.write(example.SerializeToString())
    
  2. 一旦数据集被转换为 TFRecord 格式,我们需要创建训练脚本。它将主要遵循一个典型的 TensorFlow 训练脚本,唯一的区别是我们将使用PipeModeDataset而不是TFRecordDataset。你可以使用以下代码来配置PipeModeDataset

    def _input(epochs, batch_size, channel, channel_name):
        mode = args.data_config[channel_name]["TrainingInputMode"]
        dataset = PipeModeDataset(channel=channel_name, record_format="TFRecord")
        dataset = dataset.repeat()
        dataset = dataset.prefetch(10)
        dataset = dataset.map(_dataset_parser, num_parallel_calls=10)
        if channel_name == "train":
            buffer_size = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size
            dataset = dataset.shuffle(buffer_size=buffer_size)
        dataset = dataset.batch(batch_size, drop_remainder=True)
        iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)
        image_batch, label_batch = iterator.get_next()
        return {INPUT_TENSOR_NAME: image_batch}, label_batch
    
  3. 配置 SageMaker 训练作业时,我们需要明确指定要使用 Pipe 模式:

    from sagemaker.tensorflow import TensorFlow
    hyperparameters = {"epochs": 10, "batch-size": 256}
    estimator = TensorFlow(
        entry_point="train.py",
        source_dir="3_sources",
        metric_definitions=metric_definitions,
        hyperparameters=hyperparameters,
        role=role,
        framework_version="1.15.2",
        py_version="py3",
        train_instance_count=1,
        input_mode="Pipe",
        train_instance_type="ml.p2.xlarge",
        base_job_name="cifar100-tf",
    )
    

请注意,由于 CIFAR100 数据集相对较小,你可能无法看到训练开始时间的明显减少。然而,像 COCO2017 这样的较大数据集,训练时间至少会减少几分钟。

Amazon S3 插件为 PyTorch 提供了支持

Amazon S3 插件为 PyTorch 提供了一个功能,允许你通过最少的更改直接从 S3 对象流式传输数据到现有的 PyTorch 训练脚本中。在底层,S3 插件使用 AWS SDK for C++中的TransferManager从 S3 获取文件,并利用 S3 的分段下载功能来优化数据吞吐量和可靠性。

S3 插件提供了两种 PyTorch 数据集 API 的实现:Map 风格的S3Dataset和 Iterable 风格的S3IterableDataset。在接下来的部分中,我们将讨论何时使用其中之一。

Map 风格的 S3Dataset

S3Dataset表示索引和数据记录的映射,并实现了__getitem__()方法。它允许你根据索引随机访问数据记录。当每个文件包含一个数据记录时,Map 风格的数据集效果最佳。你可以使用 PyTorch 的分布式采样器进一步将数据集在训练节点之间进行划分。

这是使用S3Dataset访问存储在 S3 上的图像的一个例子:

  1. 首先,我们将定义一个继承自父类S3Dataset的数据集类。然后,我们将使用 PyTorch 函数定义数据处理管道:

    from awsio.python.lib.io.s3.s3dataset import S3Dataset
    from torch.utils.data import DataLoader
    from torchvision import transforms
    from PIL import Image
    import io
    class S3ImageSet(S3Dataset):
        def __init__(self, urls, transform=None):
            super().__init__(urls)
            self.transform = transform
        def __getitem__(self, idx):
            img_name, img = super(S3ImageSet, self).__getitem__(idx)
            # Convert bytes object to image
            img = Image.open(io.BytesIO(img)).convert('RGB')
    
            # Apply preprocessing functions on data
            if self.transform is not None:
                img = self.transform(img)
            return img
    batch_size = 32
    preproc = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        transforms.Resize((100, 100))
    ])
    
  2. 接下来,我们将创建一个 PyTorch 原生的Dataloader对象,可以将其传递给任何训练脚本:

    # urls can be S3 prefix containing images or list of all individual S3 images
    urls = 's3://path/to/s3_prefix/'
    dataset = S3ImageSet(urls, transform=preproc)
    dataloader = DataLoader(dataset,
            batch_size=batch_size,
            num_workers=64)
    

Iterable 风格的 S3IterableDataset

S3IterableDataset表示可迭代对象,并实现了 Python 的__iter__()方法。通常,当随机读取(例如在映射风格数据集中的读取)代价高昂或不可能时,你会使用可迭代风格的数据集。当你有一批数据记录存储在单个文件对象中时,应该使用可迭代风格的数据集。

在使用S3IterableDataset时,重要的是要控制文件的大小。如果你的数据集由大量文件表示,访问每个文件都会带来开销。在这种情况下,建议将数据记录合并成更大的文件对象。

S3IterableDataset不限制可以使用的文件类型。返回的是文件对象的完整二进制数据块,你需要提供解析逻辑。你可以通过将shuffle_urls标志设置为 true 来打乱文件对象的 URL。注意,如果你需要在同一数据对象内打乱记录,可以使用ShuffleDataset,它会跨多个文件对象汇总数据记录,并从中返回一个随机样本。

S3IterableDataset会在进行分布式训练时处理数据的分片问题。你可以将S3IterableDataset与 PyTorch 的DataLoader结合使用,以实现并行数据加载和预处理。

让我们来看一个例子,如何从存储在 S3 上的多个 TAR 档案构建一个可迭代风格的数据集并应用数据转换:

  1. 我们将从使用 PyTorch 的原生IterableDataset定义一个自定义数据集类开始。在类定义的过程中,我们使用S3IterableDataset从 S3 获取数据,并应用到单个数据记录的转换:

    from torch.utils.data import IterableDataset
    from awsio.python.lib.io.s3.s3dataset import S3IterableDataset
    from PIL import Image
    import io
    import numpy as np
    from torchvision import transforms
    class ImageS3(IterableDataset):
        def __init__(self, urls, shuffle_urls=False, transform=None):
            self.s3_iter_dataset = S3IterableDataset(urls,
                                       shuffle_urls)
            self.transform = transform
        def data_generator(self):
            try:
                while True:
                    label_fname, label_fobj =       next(self.s3_iter_dataset_iterator)
                    image_fname, image_fobj = next(self.s3_iter_dataset_iterator)
                    label = int(label_fobj)
                    image_np = Image.open(io.BytesIO(image_fobj)).convert('RGB')                
                    # Apply torch vision transforms if provided
                    if self.transform is not None:
                        image_np = self.transform(image_np)
                    yield image_np, label
            except StopIteration:
                return
        def __iter__(self):
            self.s3_iter_dataset_iterator = iter(self.s3_iter_dataset)
            return self.data_generator()        
        def set_epoch(self, epoch):
            self.s3_iter_dataset.set_epoch(epoch)
    
  2. 接下来,我们定义一个转换操作来对图像进行归一化,然后实例化一个数据集实例,具备从 S3 流式传输图像的能力:

    # urls can be a S3 prefix containing all the shards or a list of S3 paths for all the shards 
     urls = ["s3://path/to/file1.tar", "s3://path/to/file2.tar"]
    # Example Torchvision transforms to apply on data    
    preproc = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        transforms.Resize((100, 100))
    ])
    dataset = ImageS3(urls, transform=preproc)
    

现在,让我们来看看 FastFile 模式。

FastFile 模式

2021 年底,亚马逊宣布了一种新的从 S3 直接流式传输数据的方法,称为 FastFile 模式。它结合了从 S3 流式传输数据的优势以及操作本地文件的便利性。在FastFile模式下,每个文件会作为 POSIX 文件系统挂载出现在你的训练程序中。因此,它与任何其他本地文件(例如存储在挂载的 EBS 卷上的文件)没有区别。

在 FastFile 模式下读取文件对象时,如果文件格式支持分块,SageMaker 将检索文件的部分内容;否则,将检索完整的文件。如果数据按顺序读取,FastFile 模式的性能最佳。请注意,检索每个文件对象会有额外的开销。因此,文件数量较少通常会导致训练作业的启动时间较短。

与之前讨论的特定于框架的流式插件相比,FastFile 模式有几个优势:

  • 避免了任何特定于框架的数据流式传输实现。你可以使用 PyTorch 或 TensorFlow 的原生数据工具,并在框架间共享数据集。

  • 因此,你可以通过使用框架实用工具对数据输入进行更精细的控制,执行诸如打乱、动态增强和实时数据处理等操作。

  • 文件格式没有限制。

  • 调试你的训练程序会更容易,因为你可以使用 SageMaker 本地模式先在本地测试和调试程序。

要使用 FastFile 模式,在配置 SageMaker 的Estimator对象时,你需要提供一个合适的input_mode值。以下代码展示了一个 TensorFlow 训练作业的例子:

from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
    entry_point="train.py",
    source_dir="3_sources",
    metric_definitions=metric_definitions,
    hyperparameters=hyperparameters,
    role=role,
    framework_version="1.15.2",
    py_version="py3",
    train_instance_count=1,
    input_mode="FastFile",
    train_instance_type="ml.p2.xlarge",
)

由于 FastFile 模式易于使用且具有多功能性,因此可以作为一个很好的起点。如果出于某种原因,你对它的性能不满意,你始终可以考虑调整数据集的配置(如文件格式、文件大小、数据处理管道、并行度等)或重新实现框架特定的实现方式。你也可以比较 FastFile 模式与其他方法(如 Pipe 模式和 PyTorch 的 S3 插件)从 S3 流式传输数据的性能。

总结

本章中,我们回顾了可用于存储和管理深度学习数据集的存储解决方案,并详细讨论了它们的优缺点及使用场景。我们通过多个示例展示了如何将 SageMaker 训练脚本与不同的存储服务集成。随后,我们了解了存储数据的各种优化策略,并讨论了优化训练任务数据检索的高级机制。我们还探讨了 SageMaker 处理服务,以及如何有效地扩展数据处理。

本章结束了本书的第一部分,该部分介绍了如何在 SageMaker 上使用深度学习模型。现在,我们将进入高级主题。在下一章中,我们将讨论 SageMaker 提供的高级训练功能。

第二部分:构建和训练深度学习模型

在这一部分,我们将学习如何使用 SageMaker 管理的功能训练深度学习模型,概述了可用的软件框架,以便在多个节点之间分布训练过程,优化硬件使用,并监控你的训练作业。

本节包含以下章节:

  • 第五章考虑硬件在深度学习训练中的作用

  • 第六章分布式训练的工程化

  • 第七章深度学习训练的操作化

第五章:考虑深度学习训练的硬件配置

训练一个大型的 深度学习DL)模型通常是一个漫长且资源消耗巨大的过程。以极端案例 GPT-3 NLP 模型为例,使用 1,024 个 NVIDIA A100 GPU 从零开始训练它大约花费了 34 天。虽然你不太可能需要从零开始训练如此庞大的模型,但即使是在自定义数据上微调大型深度学习模型,也可能需要数天甚至数周的时间。

为你的特定模型选择计算实例类型是一个至关重要的步骤,它将影响训练的成本和时间。AWS 提供了多种计算实例,适用于不同的工作负载需求。在本章中,我们将考虑最适合深度学习模型的实例的性价比特性,以及在不同场景下如何选择最合适的实例以获得最佳性能。

训练大型模型还需要将训练任务跨多个 GPU 设备和计算实例进行扩展,这一过程称为分布式训练。从高层次来看,分布式训练过程分为两个阶段:计算阶段和通信阶段。在通信阶段,单个设备和节点交换各自的更新并计算平均权重更新。交换的数据量由模型大小乘以其特性(如精度)决定。对于大型模型,训练过程中的瓶颈通常是网络吞吐量,而非单个设备的计算。因此,作为硬件考虑的一部分,我们将讨论网络吞吐量的要求以及可用的选项,如 AWS 弹性网络适配器EFA),以解决训练任务通信阶段可能出现的瓶颈。

使你的训练过程更高效的另一种方法是为特定硬件平台优化你的模型。在使用 TensorFlow 和 PyTorch 等框架训练深度学习模型时,我们依赖这些框架将模型的 Python 代码转换为要在加速器上运行的指令。然而,这些计算指令是通用的,并没有利用你的训练循环和模型架构的特定细节。SageMaker 训练编译器提供了一组功能,帮助你为特定的加速器设备优化模型,从而提高训练速度并减少内存占用。

本章将涵盖以下主题:

  • 选择最佳的计算实例

  • 使用 EFA 提升网络吞吐量

  • 使用训练编译器为 GPU 设备编译模型

阅读本章后,你将能够为你的训练任务选择高效的硬件配置,具备最佳的性价比,并进行进一步的优化。

技术要求

要跟随本章中的代码,你需要以下内容:

  • 一个 AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户

  • 已建立一个 SageMaker Notebook、SageMaker Studio Notebook 或本地兼容的 SageMaker 环境

选择最佳计算实例

Amazon SageMaker 为开发者提供了多种计算实例,按实例系列组织。每个实例系列都有一组称为实例类型的实例配置。

以下列表列出了在 SageMaker 上可用的实例系列:

  • ML.M 是一系列标准实例,提供平衡的 CPU 和内存资源配置。CPU 核心越多,内存也就越多。该实例系列不配备 GPU 设备。

  • ML.C 是一系列针对计算密集型应用(如数据处理或某些机器学习ML)算法,举例来说,支持向量机)设计的计算优化实例。该系列也可用于机器学习推理。它不配备 GPU 设备。

  • ML.G 是一系列基于 NVIDIA GPU 设备的实例,主要用于深度学习推理工作负载。它也可用于较小的训练任务和其他计算密集型工作负载。

  • ML.P 是一系列配备 NVIDIA GPU 设备的实例,专为重型深度学习训练任务设计。

到目前为止,我们只讨论了原则上可以运行任何计算操作的一般用途实例系列。除此之外,还有专门为深度学习工作负载设计的计算实例(业界称为应用特定集成电路ASIC)。在撰写本文时,有几种类型的 ASIC 实例系列可在 SageMaker 或作为 EC2 实例使用:

  • Inferentia 实例。

  • Tranium 实例,目前仅在 EC2 上可用。

  • DL1 实例,目前仅在 EC2 上可用。但已经宣布支持 SageMaker。

虽然基于 CPU 的实例可以用于运行一些机器学习训练,但它通常不是深度学习模型训练的理想选择。现在,让我们详细回顾一下可用的 GPU 和 ASIC 实例。

回顾专业的深度学习硬件

本章将重点讨论两种用于高强度深度学习训练工作负载的硬件——GPU 和 ASIC——并讨论它们为何适合深度学习训练、它们的特点以及它们的使用案例。

如果我们观察机器学习和深度学习的整体趋势,可以看到行业正从更通用的计算转向更专用的设备。

初始的机器学习和深度学习模型使用 CPU 设备进行训练,因为 CPU 允许执行几乎任何类型的计算操作。CPU 在执行单个小计算操作时也是一个优化延迟的设备。然而,大多数深度学习模型需要并行执行大量的计算操作(例如,矩阵乘法)。因此,CPU 需要花费大量时间逐一执行原子操作。

GPU 设备旨在解决不同类型的问题——并行运行大量操作。你可以说 GPU 是一种通过并行运行许多操作来优化吞吐量的设备。由于深度学习(DL)模型包括大量可以高效并行化的矩阵运算,GPU 比 CPU 更具效率。

GPU 的进步使得全新的深度学习模型架构成为可能。例如,突破性的 AlexNet 模型在 2012 年使用 GPU 设备在 ImageNet 数据集上进行了训练。研究团队将卷积和矩阵操作专门实现为在 GPU 上运行,从而在训练时实现了显著的加速。

为了简化 GPU 设备在机器学习工作负载中的使用,硬件供应商提供了用于 GPU 开发的专门库。例如,NVIDIA 创建了 CUDA 平台——一组库以及一个运行时,用于在 GPU 设备上执行通用计算。CuBLAS 库(CUDA 的一部分)提供了多种计算操作(如矩阵操作)。你还可以使用 CUTLASS 组件开发自己的操作。这对于新模型架构尤其有用。在 CUDA 上优化计算操作也能提升训练性能。

最近,一种新的深度学习硬件设计方法变得流行起来:ASIC。这是一种旨在执行有限集操作的设备,但能极其高效地完成这些操作。谷歌的张量处理单元TPU)就是一种为深度学习工作负载设计的 ASIC 示例。AWS 也在积极开发用于深度学习工作负载的专用硬件设备。目前,AWS 已推出 Inferentia(2018 年)用于推理,Tranium(2021 年)和 DL1 实例(2022 年)基于 Gaudi 加速器用于训练。请注意,在撰写本文时,Tranium 和 DL1 加速器仅作为 EC2 实例提供。我们预计它们未来将在 SageMaker 上可用。

由于 ASIC 的高度专业化,确认特定的深度学习框架或模型架构是否支持某个 ASIC 设备始终是个好主意。通常,你需要将模型代码转换为 ASIC 的指令。这通常由提供的编译器自动完成。对于 AWS 的 ASIC,你需要使用开源的 Neuron SDK 编译你的模型(aws.amazon.com/machine-learning/neuron/)。

在编译你的模型时,Neuron SDK 提供了多种优化方法,例如将操作进行批处理。它使用提前编译,因此输入数据批次的维度应该在模型配置时提前定义,尽管需要注意的是,Neuron SDK 还支持一套已定义的操作符。如果你的模型有不受支持的操作符(例如自定义控制流操作),你将无法编译你的模型。Neuron SDK 支持 TensorFlow、PyTorch 和 MXNet 框架。

在许多情况下,选择最佳的 ASIC 或 GPU 设备取决于你的特定模型和训练超参数。你可以使用业界标准基准测试 MLPerf(mlcommons.org/en/training-normal-11/)作为参考。领先的 GPU 和 ASIC 厂商在八个流行的深度学习模型上训练后,提交了其硬件加速器的性能详情。截至 2021 年 12 月,NVIDIA A100 GPU 在所有模型上表现优于市面上所有可用的硬件加速器。Google 的 TPUv4 ASIC 加速器在六个模型上提高了基准测试成绩,但在提交时,TPUv4 尚未上市。

选择最佳实例类型

你选择的实例系列和具体实例类型始终由你的使用案例需求驱动。重要的是,你可以在同一使用案例中使用多个实例类型和系列。例如,你可能想先从单 GPU 训练开始,同时调整超参数并进行整体模型调试。然后,你可以逐步将训练扩展到更多节点,或将训练任务迁移到具有更高性能加速器的实例类型上。

在选择最佳实例类型时,你必须考虑以下一些标准:

  • 模型架构及其大小:这决定了在 GPU 加速器上存储模型的内存需求。

  • 所需训练模式:在这里,你需要选择是否希望在单个 GPU、多个 GPU 或多个 GPU 多节点上进行训练。

  • 业务优先级:在这里,你需要选择是想尽可能快地训练模型,还是尽可能便宜,或者找到一个可接受的性价比平衡点。

在选择适合自己特定情况的实例时,牢记以下实例类型的特点是很重要的:

  • 加速器架构:这会影响计算性能。例如,最新的 NVIDIA A100 芯片比上一代 V100 芯片提升了约 2.5 倍的性能。

  • 可用 vCPU 核心:这些将用于数据加载和处理等操作。

  • GPU 内部和节点间的网络吞吐量:这定义了在进行多 GPU 和/或多节点训练任务时,数据(梯度)在训练设备之间交换的速度。

  • 选择的实例类型的价格。

在以下小节中,我们将概述几个典型的使用案例,从最小且最具成本效益的模型到最大且最具性能的模型按顺序排列。

G5 系列——为小型和中型模型提供高效的训练

在实验中训练小型或中型深度学习模型时,你可以考虑使用 G5 实例,因为它们具有成本效益且性能强大。它们配备最多八个NVIDIA A10G加速器,最高 100 Gbps 网络带宽,以及最多 192 个 vCPU。下表展示了 G5 系列的规格:

实例规格 GPU GPU 内存(GiB) vCPUs 内存(GiB) 网络带宽(Gbps)
单 GPU 虚拟机 g5.xlarge 1 24 4 16
g5.2xlarge 1 24 8 32 高达 10
g5.4xlarge 1 24 16 64 高达 25
g5.8xlarge 1 24 32 128 25
g5.16xlarge 1 24 64 256 25
多 GPU 虚拟机 g5.12xlarge 4 96 48 192
g5.24xlarge 4 96 96 384 50
g5.48xlarge 8 192 192 768 100

图 5.1 – G5 系列规格

如果你希望在单个 GPU 设备上运行模型,你应该根据其他系统需求(网络、RAM、vCPUs 等)选择单 GPU 虚拟机。如果你希望同时运行多个实验(每个使用不同的 GPU 设备),你应该选择多 GPU 虚拟机。需要注意的是,在多 GPU 虚拟机的情况下,单个 GPU 设备之间并未通过高速的NVLink 互联连接。因此,如果你打算进行多 GPU 分布式训练,带有 NVLink 的 P3 系列会更加合适。

另外,你也可以考虑上一代的 G4 实例,它的小时价格较低(某些实例的价格最高可比 G5 便宜 50%)。然而,根据 AWS 内部基准测试,G5 在性价比方面比 G4 高出多达 40%。

P3 系列 – 高性能与高性价比的训练

P3 系列提供高性能和高性价比,适用于大规模模型。它最多配备八个NVIDIA V100加速器,并且与 G5 系列不同,它支持高效的 NVLink GPU 互联:

实例规格 GPU – Tesla V100 GPU 点对点 GPU 内存(GB) vCPUs 内存(GB) 网络带宽
p3.2xlarge 1 16 8 61 高达 10 Gbps
p3.8xlarge 4 NVLink 64 32 244 10 Gbps
p3.16xlarge 8 NVLink 128 64 488 25 Gbps
p3dn.24xlarge 8 NVLink 256 96 768 100 Gbps

图 5.2 – P3 系列规格

p3.2xlarge实例是运行复杂 DL 模型单 GPU 训练的不错选择(假设模型可以装入内存)。如果你的模型无法适应单个 GPU 设备,你可以选择p3.8xlargep3.16xlarge,它们是多节点实例。在这种情况下,你将把模型的部分存储在多个 GPU 中。NVLink 互联在前向和反向传播过程中提供 GPU 之间的高速数据交换。

p3.8xlargep3.16xlarge的另一个应用领域是运行多 GPU 数据并行训练任务。在这种使用场景中,你将深度学习模型的副本加载到每个 GPU 设备中,但使用不同的数据批次进行训练。NVLink 互联确保在每次训练迭代结束时,GPU 节点之间进行高速的梯度交换和计算。

最强大的实例p3dn.24xlarge配备了 EFA 网络设备,提供低延迟和一致的节点间通信。这使得p3dn.24xlarge实例成为大规模多 GPU 多模式训练任务的优秀选择,尤其是当你的训练任务受限于网络时。

P4 系列 – 训练的最高性能

P4 系列基于 NVIDIA A100 GPU 加速器,截至 2021 年 12 月,超越了任何市售加速器在 MLPerf 基准测试中的表现。P4 系列有一个单一实例,p4d.24xlarge,配备了八个 A100 GPU 设备、96 个 vCPU 和 1,152 GB 的 RAM。

这些特性使得p4d.24xlarge实例成为使用分布式训练方法训练大型 SOTA 深度学习模型的理想选择。然而,在训练大型模型时,训练集群中设备之间需要交换的数据量可能会超过 GPU 之间和节点之间的网络带宽,这可能导致整体训练速度变慢,并且浪费昂贵的 GPU 资源。AWS 为p4d.24xlarge实例提供了几种网络功能来缓解此问题:

  • 通过 NVLink,在同一节点内的 GPU 之间实现高达 600 GB/s 的双向带宽

  • 通过 EFA 使用GPUDirect RDMA,在不同节点之间实现高达 400 GB/s 的带宽

此外,p4d.24xlarge支持广泛的精度点类型:FP64、FP32、FP16、INT8、BF16 和 TF32。如果你的框架和模型支持混合精度,你可能能够在不大幅牺牲模型精度的情况下实现更好的性能。

自然地,p4d.24xlarge的价格高于其他实例。然而,第二昂贵的实例p3dn.24xlarge与之相比的价格差距只有约 5%。根据其卓越的性能,P4 可以为训练提供高达 60%的成本节约,并且根据 AWS 内部基准测试,深度学习性能提高超过 2.5 倍。这使得p4d.24xlarge不仅是深度学习训练中性能最强的实例,也是训练大型 SOTA 深度学习模型中性价比最高的选择。你可以在以下文章中找到有关 P4d 实例系列的详细性能基准: aws.amazon.com/blogs/compute/amazon-ec2-p4d-instances-deep-dive/

通过 EFA 提高网络吞吐量

在训练大型深度学习模型时,你需要将大型训练任务分解成更小的任务,并分布到多个计算设备上。分布式训练包括以下关键步骤:

  1. 训练集群中的每个设备执行以下操作:

    1. 从全局数据批次中读取一个独特的 mini-batch

    2. 通过模型运行一个 mini-batch 并计算损失

    3. 计算梯度以最小化损失

  2. 每个设备将梯度传递给其同伴。计算平均梯度。

  3. 每个设备根据平均梯度更新模型。

为了衡量分布式训练的效率,我们可以使用扩展因子,定义如下:

这里,T 是单个设备的吞吐量,n 是训练集群中的设备数量,nT 是你的训练集群所实现的整体吞吐量。虽然理想的扩展性通常难以实现(即增加资源并不总是能成比例地减少训练时间),但在许多最近的基准测试中,已经证明通过精心应用、硬件和网络优化,可以实现高达 90% 的扩展效率。

为了对你的训练任务进行性能瓶颈分析,测量每个步骤的性能非常重要。研究表明,在许多情况下,通信阶段(步骤 2)是训练过程中的全球性瓶颈(例如,可以参考 网络是分布式训练瓶颈吗?arxiv.org/pdf/2006.10103.pdf)。在本节中,我们将重点理解如何优化通信阶段。

有几个因素决定了每个节点发送和接收的数据量:

  • 首先是通信算法,它定义了训练设备之间如何交换梯度更新。在撰写时,最流行的方法被称为 Ring-AllReduce。这个算法使你能够高效地在每个训练设备之间传输梯度更新。每个 N 节点与其两个同行进行通信 次。每个训练设备在单次迭代中发送的总体信息量是 ,对于大 N 来说,其中 D 是梯度更新的大小。以下图所示:

图 5.3 – Ring-AllReduce 通信算法

图 5.3 – Ring-AllReduce 通信算法

  • 其次是模型的大小及其精度(在前面的公式中为 D)。

例如,如果我们使用 Ring-AllReduce 算法来训练 BERT 模型(该模型包含大约 3.4 亿个参数)并采用半精度,每个训练设备在单次迭代中将发送和接收大约 650 MB 的数据。通信需要快速进行。单个设备的性能下降将导致整体训练过程的放缓。

引入 EFA

Amazon EFA 是一种网络设备,提供比传统 TCP 传输更低且更稳定的延迟。EFA 专门为高性能和机器学习应用场景设计,其中实例间通信对分布式任务至关重要。

EFA 提供了以下优势:

  • 操作系统绕过功能,允许深度学习应用直接与网络接口硬件通信,从而提供低延迟和可靠的传输功能。

  • 支持高性能消息协议,如MPINCCL。对于深度学习用例,我们特别关注 NVIDIA 的 NCCL 库,它为 GPU 设备提供高性能的通信例程。

使用 EFA 可以显著提高训练任务的性能。根据 AWS 基准测试,使用 EFA 可以使 BERT 在 ml.p4dl.24xlarge 的 32 个实例上训练速度比默认的 弹性网络适配器 (ENA) 快 130%。

在 SageMaker 上使用 EFA 不会产生额外费用。EFA 可用于 ml.p3dn.24xlargeml.p4d.24xlargeml.c5n.18xlarge SageMaker 实例。

在自定义训练容器中使用 EFA

SageMaker 提供与 EFA 设备的无缝集成。如果您使用 TensorFlow 或 PyTorch 深度学习容器与支持的训练实例,EFA 将自动启用。

如果您选择使用自定义容器,您需要在该容器中安装必要的 EFA 包以及 MPI 和 NCCL 库。以下步骤展示了如何在 Dockerfile 中执行这些操作:

  1. 首先,您需要定义您将使用的 MPI、NCCL、EFA 和 OFI 库的版本,如下所示:

    ARG OPEN_MPI_PATH=/opt/amazon/openmpi/
    ENV NCCL_VERSION=2.7.8
    ENV EFA_VERSION=1.11.2
    ENV BRANCH_OFI=1.1.1
    
  2. 然后,您必须下载并执行 EFA 驱动程序安装程序:

    RUN cd $HOME \
      && curl -O https://s3-us-west-2.amazonaws.com/aws-efa-installer/aws-efa-installer-${EFA_VERSION}.tar.gz \
      && tar -xf aws-efa-installer-${EFA_VERSION}.tar.gz \
      && cd aws-efa-installer \
      && ./efa_installer.sh -y --skip-kmod -g \
    ENV PATH="$OPEN_MPI_PATH/bin:$PATH"
    ENV LD_LIBRARY_PATH="$OPEN_MPI_PATH/lib/:$LD_LIBRARY_PATH"
    
  3. 现在,您必须从公共的 NVIDIA 仓库克隆并构建 NCCL 库:

    RUN cd $HOME \
      && git clone https://github.com/NVIDIA/nccl.git -b v${NCCL_VERSION}-1 \
      && cd nccl \
      && make -j64 src.build BUILDDIR=/usr/local
    
  4. 接下来,您必须安装 AWS OFI NCCL 插件,它允许您将 EFA 网络模块与 NCCL 应用程序一起使用:

    RUN apt-get update && apt-get install -y autoconf
    RUN cd $HOME \
      && git clone https://github.com/aws/aws-ofi-nccl.git -b v${BRANCH_OFI} \
      && cd aws-ofi-nccl \
      && ./autogen.sh \
      && ./configure --with-libfabric=/opt/amazon/efa \
           --with-mpi=/opt/amazon/openmpi \
           --with-cuda=/usr/local/cuda \
           --with-nccl=/usr/local --prefix=/usr/local \
      && make && make install
    
  5. 最后,您必须安装 NCCL 测试并执行它们,以检查 NCCL 操作的正确性和性能:

    RUN cd $HOME \
      && git clone https://github.com/NVIDIA/nccl-tests \
      && cd nccl-tests \
      && make MPI=1 MPI_HOME=/opt/amazon/openmpi CUDA_HOME=/usr/local/cuda NCCL_HOME=/usr/local
    

在本节中,我们讨论了分布式训练中设备之间的网络以及它对整体训练效率的影响。由于网络经常成为训练的全球瓶颈,我们分享了如何根据集群配置和模型参数来估算网络带宽。然后,我们回顾了 AWS 的 EFA 网络设备,它提高了网络带宽和效率。由于 EFA 对用户没有额外费用或任何缺点,因此建议在可能的情况下使用它。

使用训练编译器为 GPU 设备编译模型

SageMaker 训练编译器是一项能力,允许您自动优化 NLP 深度学习模型,以便在 GPU 实例上运行。对于支持的模型架构和框架,您的训练脚本无需进行代码更改。您只需在 SageMaker 训练任务配置中启用训练编译器即可。训练编译器不仅可以减少训练时间和内存需求,而且不会影响模型的准确性。例如,根据 AWS 基准测试,使用训练编译器时,基于 RoBERTa 的模型训练时间和成本降低了 30%。

让我们回顾一下 SageMaker 训练编译器的工作原理,以及如何在训练任务中使用它。

引入 XLA 优化库

加速线性代数XLA)是一种特定领域的编译器,它通过几乎不修改模型代码的方式加速模型训练和执行。到目前为止,XLA 已经支持 TensorFlow 和 PyTorch 框架。SageMaker 训练编译器抽象化了与 XLA 库的交互,并利用它们来优化在 SageMaker 上运行的训练任务。SageMaker 训练编译器支持单 GPU 和分布式训练任务。

当你在没有 XLA 的情况下训练模型时,所有操作都会单独执行。假设你的模型有两个顺序操作:矩阵乘法和矩阵加法。没有 XLA 时,你的框架执行引擎会将这两个操作(称为 kernels)一个接一个地发送到 GPU 设备。当使用 XLA 时,它会通过将加法和乘法操作融合,将这两个操作编译成一个单独的内核启动。融合的操作必须完全在 GPU 寄存器上执行,只有最终的结果会被流式传输给用户。去除冗余的内存操作是 XLA 编译器的关键优化特性之一。

XLA 编译器与其他编译器的另一个显著区别是,不像常规的 CUDA 操作会立即执行(称为 eager 执行),XLA 张量操作是“惰性”的。首先,XLA 编译器构建融合操作的图,并将张量作为占位符保留在这个执行图中。只有在操作结果需要时,计算操作才会被执行。通过延迟执行,XLA 能够在模型的计算图中找到融合操作的机会。

使用 SageMaker 训练编译器

SageMaker 训练编译器已在广泛的 NLP 模型上进行了测试,并且也支持流行的计算机视觉模型,如用于图像分类和目标检测的 PyTorch 和 TensorFlow 实现。随着我们预计这个列表会随着时间增长,请查阅以下页面以了解最新的支持模型:docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-support.xhtml。此页面还提供了建议的训练和模型配置,如实例类型、精度(是否混合)和批大小。

SageMaker 训练编译器也可以用于尚未正式测试的模型。在使用训练编译器处理未测试的模型时,请注意以下事项:

  • 你可能需要修改你的训练脚本,例如设置合适的 XLA 设备,使用兼容 XLA 的优化器、数据加载器和 XLA 训练循环语义。

  • 你可能需要进行超参数搜索(特别是批大小和学习率),以找到适合你训练任务的最优配置。这是因为 SageMaker 训练编译器会改变你模型的内存占用。

  • 训练编译器仅适用于部分 SageMaker 深度学习容器。请参阅以下页面以获取最新支持训练编译器的容器:github.com/aws/deep-learning-containers/blob/master/available_images.md

在基准测试您的自定义模型时,比较启用和未启用训练编译器的结果,请记住,SageMaker 编译您的模型需要一些时间,这会增加整体训练时间。因此,可能不适合在短时间训练作业(例如在小数据集上微调任务)中使用训练编译器。此外,正确设置批处理大小也很重要。通常,您可以期待训练编译器减少模型的内存占用,从而增加最大批处理大小。随着批处理大小的增加,您需要按比例调整学习率。请注意,给定模型的内存需求可能并不总是会减少。在这种情况下,您将无法增加批处理大小。对于未经测试的模型,使用训练编译器需要通过实验来获得最佳结果。

使用训练编译器

若要在已测试的模型中使用训练编译器,您需要显式地在训练作业配置中启用它。请按照以下步骤操作:

  1. 首先导入 TrainingCompilerConfig 对象。请注意,它适用于 PythonSDK > 2.7.x:

    from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig
    

TrainingCompilerConfig 对象支持以下参数:

  • enabled (bool):可选项。此项是一个开关,用于启用 SageMaker 训练编译器。默认值为 True

  • debug (bool):可选项。此项指定是否输出调试的详细日志,这可能会导致性能下降。默认值为 False

  1. 接下来,您需要为 SageMaker 训练编译器配置必要的超参数:

    hyperparameters = {
        "epochs": 5,
        "train_batch_size": 24,
        "model_name": "bert-base-cased",
    }
    # Scale the learning rate by batch size, as original LR was using batch size of 32
    hyperparameters["learning_rate"] = float("5e-5") / 32 * hyperparameters["train_batch_size"]
    
  2. 接下来,您必须像之前一样配置 HuggingFace 训练作业,唯一的区别是您必须显式传递 TrainingCompilerObject,并将其设置为训练配置中的默认 enabled 状态:

    sm_training_compiler_estimator = HuggingFace(
        entry_point="train.py",
        instance_type="ml.p3.2xlarge",
        instance_count=1,
        role=role,
        py_version="py38",
        transformers_version="4.11.0",
        pytorch_version="1.9.0",
        compiler_config=TrainingCompilerConfig(),
        hyperparameters=hyperparameters,
        disable_profiler=True,
        debugger_hook_config=False,
    )
    

注意

为了获得训练编译器的最佳性能,建议禁用 SageMaker 配置文件和 SageMaker 调试器功能。请注意在我们的训练作业中设置适当的配置。

一旦训练作业启动,您必须确保模型已被编译。为此,您应期望在训练作业日志中看到以下消息,这表示训练编译器按预期工作:

Found configuration for Training Compiler
Configuring SM Training Compiler...

现在,让我们总结一下本章内容。

总结

在本章中,我们重点讨论了工程化深度学习分布式训练的硬件方面。我们回顾了可用的 SageMaker 计算实例,并聚焦于带有 GPU 设备的实例系列。接着,我们讨论了不同的深度学习用例,以及如何为它们选择最优的计算实例。然后,我们回顾了分布式训练的网络需求,并学习了 Amazon EFA 如何帮助你避免在运行大规模训练任务时遇到网络瓶颈。我们还回顾了如何使用 SageMaker 训练编译器优化模型,使其能够在 GPU 设备上运行,并通过实践体验了这一功能。

在下一章,第六章分布式训练工程化,我们将继续讨论分布式训练。我们将重点讨论如何为你的用例、深度学习框架和模型架构选择最合适的分布式训练类型,并在这些领域积累实践经验。

第六章:分布式训练工程

在上一章中,我们讨论了如何为深度 学习DL)训练任务选择最优硬件,并针对目标硬件平台优化你的模型。在本章中,我们将深入探讨如何根据你的特定用例和模型架构,在亚马逊 SageMaker 上设计高效的分布式训练。

分布式训练旨在解决两个具体问题。第一个问题是如何通过将训练任务分配到多个计算设备上来减少大型模型的训练时间。另一个问题是在需要训练无法完全加载到单个 GPU 设备内存中的大型模型时出现的。这一问题在自然语言处理(NLP)任务中尤为重要,因为研究表明,超大模型具有更强的表达能力,因此在广泛的 NLP 任务中表现更好。例如,最新的开源 SOTA 语言模型 BLOOM,在一个包含 384 个 GPU 加速器(NVIDIA A100)的计算集群上训练了大约 3.5 个月。仅模型权重就有约 329 GB,而包含模型权重和优化器状态的检查点则为 2.3 TB。更多详情,请参阅 huggingface.co/bigscience/bloom

为解决这些问题,已经出现了两种方法;第一种是数据并行分布式训练,通过同时分配任务来加速训练时间。第二种是模型并行分布式训练,将大型模型分布到多个 GPU 之间,从而使你能够使用无法完全加载到单个 GPU 设备内存中的模型。

正如你可能已经猜到的那样,无法适配单个 GPU 设备的大型模型也需要相当长的训练时间。因此,必然需要将模型并行与数据并行结合起来,以使训练时间变得可接受。数据并行和模型并行的结合被称为混合并行。在本章中,我们将讨论这三种并行方式。

虽然理解分布式训练方法至关重要,但你还需要了解适用于你的 DL 框架和模型架构的实现方式。SageMaker 提供了用于分布式训练的专有库:SageMaker 分布式数据并行SDDP)和SageMaker 分布式模型并行SDMP)。在本章中,我们将回顾它们的优点,并实际体验如何使用它们。此外,我们还将讨论针对 TensorFlow 和 PyTorch 框架的其他流行开源分布式训练替代方案,以及如何在 SageMaker 平台上使用它们。

本章将涵盖以下主题:

  • 数据并行训练工程

  • 模型并行与混合并行训练工程

  • 优化分布式训练任务

到本章结束时,您将对分布式训练有一个清晰的理解,并且获得在 Amazon SageMaker 上实现各种类型分布式训练的实际经验。

技术要求

在本章节中,我们将提供代码示例,帮助您培养实际技能。完整的代码示例可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter6/查看。

要跟随本代码示例,您需要具备以下条件:

  • 一个 AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户。

  • 建立一个 SageMaker 笔记本,SageMaker Studio 笔记本,或本地兼容 SageMaker 的环境。

  • 在您的 AWS 账户中访问 GPU 训练实例。本章节中的每个示例将提供推荐的实例类型。可能需要增加您的计算配额,以便启用 GPU 实例用于SageMaker 训练作业。在这种情况下,请按照docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.xhtml中的说明操作。

工程数据并行训练

首先,让我们概述一下在本章中将使用的一些重要术语:

  • 训练进程训练器工作节点 – 这些术语交替使用,指的是计算集群中的独立训练进程。例如,分布式深度学习训练进程通常在单个 GPU 设备上运行。

  • 训练节点服务器主机 – 这些术语定义了训练集群中的服务器。该服务器可以有一个或多个 GPU 设备,这意味着一个或多个训练进程可以在同一台服务器上运行。

  • 世界大小 – 这是在训练集群中运行的独立训练进程的数量。通常,世界大小等于训练集群中可用的 GPU 设备数量。

  • 排名(也称全局排名) – 这是在训练集群中运行的训练进程的唯一零基 ID。例如,如果您有 4 个训练进程,它们的排名分别为 0、1、2 和 3。

  • 本地排名 – 这是在单个节点内运行的训练进程的唯一零基 ID。例如,如果您有两个训练节点,每个节点有两个 GPU 设备,那么本地排名将是 0 和 1,全局排名将是 0、1、2 和 3。

  • 通信后端集体通信 – 这些术语定义了训练进程之间进行通信和协调计算的机制与协议。一些常见的后端有NVIDIA NCCLGloo消息传递接口MPI)。

  • allreduce 操作用于聚合和平均张量或广播,将张量从一个训练进程发送到集群中的其他进程。通常,通信后端提供集体操作的实现。

现在我们了解了分布式训练的基本术语,让我们深入回顾一下数据并行。

数据并行分布式训练在你想要减少在多个训练设备上训练模型的时间时非常有用。每个训练进程都有一个全局模型的副本,但它在与其他进程并行的过程中,在独特的数据切片上进行训练(因此称为数据并行性)。在训练步骤结束时,每个训练进程与其他进程交换学习到的梯度更新。然后,梯度更新被平均并分发回所有训练进程,以便它们可以更新各自的模型副本。图 6.1展示了数据批次在一个数据并行的双节点双 GPU 集群中的分布情况:

图 6.1 – 数据并行概览

图 6.1 – 数据并行概览

在设计数据并行训练任务时,你需要注意几个关键设计选择,以便调试和优化你的训练任务,诸如以下几点:

  • 各个进程之间如何进行协调

  • 各个计算进程如何相互通信

  • 计算进程如何在训练集群中分布

在接下来的章节中,我们将讨论这些设计选项。

协调模式 – 参数服务器与 Allreduce

在分布式集群中协调计算进程有两种方式:使用专用的集中式协调器,以及使用对等协调,其中每个节点直接与集群中的一个或多个对等节点通信。在数据并行训练的背景下,集中式协调模式被称为参数服务器,其中参数服务器进程协调梯度更新的分发并维护全局模型副本。而对等模式被称为Allreduce,它是一种通过对等算法在训练进程之间分发梯度更新的方式。在图 6.2中,你可以看到这两种协调模式的区别:

图 6.2 – 参数服务器 (A) 和 Allreduce (B) 协调模式

图 6.2 – 参数服务器 (A) 和 Allreduce (B) 协调模式

参数服务器负责协调集群中的训练过程,具体包括以下内容:

  • 为每个训练过程分配一组独特的数据记录

  • 从每个训练进程接收梯度

  • 聚合梯度并相应地更新模型权重

  • 将更新后的模型发送回训练进程

参数服务器存储模型权重的主副本。对于较大的深度学习模型,可能无法将完整的模型存储在参数服务器上。此外,参数服务器可能会成为网络和计算瓶颈。在这种情况下,你可能需要引入多个参数服务器,它们将存储模型参数的子集,从而减少网络和内存的需求。多个参数服务器允许你扩展分布式训练以处理大型模型;然而,当在训练进程和参数服务器之间协调模型更新时,它会引入额外的复杂性,并可能导致网络拥塞。找到训练进程与参数服务器之间的最佳配置可能是一个艰巨的任务,需要经过相当多的反复试验才能找到最优配置。

Allreduce 算法采用点对点通信,每个训练进程仅与两个邻居交换梯度更新。一个秩为 i 的训练进程为独特的数据微批次计算梯度,从进程 i-1 接收梯度,并将接收到的梯度与自己计算的梯度汇总,然后将聚合后的梯度发送给节点 i+1。总的来说,每个进程将与它的同伴进行 次通信:

图 6.3 – Allreduce 算法中的计算操作顺序

图 6.3 – Allreduce 算法中的计算操作顺序

Allreduce 算法被认为是带宽高效的,具有恒定的通信成本,并避免了像参数服务器那样的通信瓶颈。此外,与参数服务器方法相比,它的操作复杂性较低(特别是在多个参数服务器实例的情况下)。因此,许多近期的研究论文和实现都基于 Allreduce 算法及其修改版。Allreduce 算法的最流行实现包括 Horovod、TensorFlow Mirror Strategy 和 PyTorch 分布式数据并行DDP)。AWS 也在 SDDP 库中使用了修改后的 Allreduce 算法。稍后在本章中,我们将使用前面提到的 Allreduce 实现,在 SageMaker 上开发一个分布式训练任务。

通信类型 – 同步与异步

分布式训练任务中有两种通信类型:同步sync)和异步async)。

同步通信意味着每个训练进程将与集群中的其他进程同步进行计算。例如,在同步 Allreduce 算法的情况下,每个训练进程将在其他进程完成其反向和正向传递后,等待交换其梯度。这会导致集群在每个训练步骤上的性能由最慢的训练进程决定,并可能导致其他训练进程的等待时间(浪费)。然而,同步通信的好处包括更稳定的训练收敛。Allreduce 算法的不同实现也提供了优化,以减少等待时间。

异步通信中,每个节点独立工作。它将梯度更新发送给其他进程或集中式参数服务器,并在不等待同伴结果的情况下继续进行下一个训练迭代。这种方法可以最小化等待时间,并最大化每个训练进程的吞吐量。该方法的缺点是,由于增加了随机性,训练过程可能会变得收敛较慢且不稳定。

实际上,平衡系统吞吐量和训练收敛是非常重要的。为此,大多数分布式训练实现都使用同步通信,并通过多种优化来提高训练吞吐量。

集群中的训练进程布局

根据你的模型大小/架构和训练需求(例如所需的训练时长),有几种方法可以在训练集群中组织训练进程:

  • p4d.24xlarge实例仅有 8 个 GPU 设备,这限制了在单节点上扩展训练任务的能力。

  • 多个节点单个 GPU - 这意味着所有进程之间的协调都通过网络通信完成,这通常会成为全球训练瓶颈。因此,这种布局对于大多数训练场景来说是次优的。

  • 多个节点多个 GPU - 这允许你将训练任务扩展到数十个甚至数百个独立的训练进程。在选择此布局时,你需要关注训练节点之间的网络吞吐量,因为它可能成为全球瓶颈。SageMaker 实例如p4dp3dn提供了改进的网络能力来解决这个问题。

现在我们已经对数据并行性有了初步的直觉,接下来让我们积累一些实践经验,并为 TensorFlow 和 PyTorch 框架构建数据并行分布式训练任务。我们将使用本地数据并行实现和深度学习框架,以及流行的与框架无关的 Horovod 库。然后,我们将学习如何使用 AWS 的专有 SageMaker 数据并行库,并与开源数据并行实现进行对比,了解其优势。

工程化 TensorFlow 数据并行训练

在为 TensorFlow 框架设计分布式训练作业时,你可以使用几种数据并行实现:

  • 原生数据并行实现(称为“策略”)

  • TensorFlow 的 Horovod 实现

让我们回顾一下这些实现的优缺点。

使用原生 TensorFlow 策略

与 TensorFlow 1 相比,TensorFlow 2 大幅扩展了分布式策略的数量。请注意,由于 TensorFlow 2 提供了多个训练 API,部分 API 对分布式策略的支持有限。请参见图 6.4中的支持矩阵:

训练 API 镜像策略 TPU 策略 多工作节点镜像策略(MWMS) 中央存储策略 参数服务器策略
Keras Model.fit 支持 支持 支持 实验性支持 实验性支持
自定义训练循环 支持 支持 支持 实验性支持 实验性支持
Estimator API 支持有限 不支持 支持有限 支持有限 支持有限

图 6.4 – TensorFlow 2 分布式策略

参数服务器策略中央存储策略被标记为实验性支持,这意味着它们目前正在积极开发中。通常建议在生产工作负载中不要使用实验性功能。因此,我们不会在本书的范围内考虑它们。

注意

虽然 Amazon SageMaker 文档声明它支持 TensorFlow 参数服务器,但这一说法具有误导性。SageMaker 支持的是 TensorFlow 1 参数服务器,它已经过时,不应在新开发中使用。SageMaker 并不直接支持 TensorFlow 2 的原生策略,尽管通过一些代码更改可以支持它们,接下来会展示这一过程。

TPU 策略旨在与 Google TPU 设备配合使用,因此不受 Amazon SageMaker 支持。因此,在本节中,我们将重点关注 镜像策略 和 MWMS。

这两种策略都实现了针对 GPU 设备的同步 Allreduce 算法。顾名思义,多工作节点策略支持将训练任务分布到多个训练节点上。对于节点内部通信,你可以选择使用NCCL 后端原生 RING 通信后端。在这两种策略中,完整的模型副本(称为镜像变量)存储在每个训练进程中,并在每次训练步骤后同步更新。让我们回顾一个在 SageMaker 平台上实现 MWMS 的示例。

作为测试任务,我们将选择大家最喜爱的 MNIST 数据集,并训练一个小型计算机视觉模型来解决分类任务。我们将使用方便的Keras API来构建并训练模型,并评估结果。更多细节的示例笔记本可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter6/1_distributed_training_TF.ipynb中找到。

我们将首先回顾启用 MWMS 所需的修改。

集群配置和设置

MWMS 在 Amazon SageMaker 中并未原生支持,因此我们需要在 SageMaker 中正确配置 MWMS 环境。TensorFlow2 使用一个名为tf_config的环境变量来表示集群配置。此配置随后用于启动训练过程。您可以在www.tensorflow.org/guide/distributed_training#TF_CONFIG了解如何构建'TF_CONFIG'变量。在以下代码块中,我们使用'_build_tf_config()'方法来设置此变量。请注意,我们在此过程中使用了'SM_HOSTS''SM_CURRENT_HOST'的 SageMaker 环境变量:

Def _build_tf_config():
    hosts = json.loads(os.getenv("SM_HOSTS"))
    current_host = os.getenv("SM_CURRENT_HOST")
    workers = hosts
    def host_addresses(hosts, port=7777):
        return ["{}:{}".format(host, port) for host in hosts]
    tf_config = {"cluster": {}, "task": {}}
    tf_config["cluster"]["worker"] = host_addresses(workers)
    tf_config["task"] = {"index": workers.index(current_host), "type": "worker"}
    os.environ["TF_CONFIG"] = json.dumps(tf_config)

在此示例中,默认情况下,我们使用两个p2.xlarge实例,训练过程的世界大小仅为两个。因此,_build_tf_config()将在 rank=0节点中生成以下'TF_CONFIG'变量:

{
    "cluster": 
    {
        "worker": ["algo-1:7777", "algo-2:7777"]},
        "task": {"index": 0, "type": "worker"
    }
}

一旦TF配置正确设置,TF2 应该能够在所有节点上启动训练过程,并利用所有可用的 GPU 设备。 这是默认设置,但您也可以提供要使用的特定 GPU 设备列表。

要完成集群设置,我们还需要确保已经配置了 NCCL 后端(请参见_set_nccl_environment()方法),并确保集群中的所有节点可以相互通信(请参见_dns_lookup()方法)。请注意,这些方法是必需的,因为 TensorFlow 2 策略并未得到 SageMaker 的官方支持。对于受支持的数据并行实现,SageMaker 提供了这些工具,并作为训练集群初始化的一部分进行运行。

使用 MWMS

为了使用 MWMS,我们将首先启动一个策略对象,如下所示。请注意,在这里,我们显式设置通信后端为AUTO,这意味着 TF2 将自动识别使用哪个后端。您也可以手动定义特定的后端。NCCL和自定义的RING后端可用于 GPU 设备:

strategy = tf.distribute.MultiWorkerMirroredStrategy(
    communication_options=tf.distribute.experimental.CommunicationOptions(
        implementation=tf.distribute.experimental.CollectiveCommunication.AUTO
    )
)

一旦策略被正确初始化,您可以通过检查strategy.num_replicas_in_sync来确认集群配置,它将返回您的世界大小。它应该与每个节点的 GPU 数量乘以节点数量匹配。

在这个示例中,我们使用的是 Keras API,它完全支持 MWMS,因此简化了我们的训练脚本。例如,要在所有工作节点上创建模型副本,你只需要在strategy.scope内初始化你的 Keras 模型,如下代码块所示:

    with strategy.scope():
        multi_worker_model = build_and_compile_cnn_model()

此外,MWMS 会自动根据世界大小分片你的数据集。你只需要设置一个合适的全局批处理大小,如下代码块所示。请注意,如果需要某些自定义分片逻辑,可以开启自动分片功能:

global_batch_size = args.batch_size_per_device * _get_world_size()
multi_worker_dataset = mnist_dataset(global_batch_size)

剩余的训练脚本与单进程 Keras 训练脚本类似。正如你所见,使用 MWMS 是相当直接的,TF2 在抽象化复杂性方面做得很好,同时,如果需要,也为你提供了调整默认设置的灵活性。

运行 SageMaker 任务

到目前为止,我们已经讨论了如何更新训练脚本以进行数据并行训练。在源目录中,你还将看到 mnist_setup.py 脚本,用于下载并配置 MNIST 数据集。现在我们已经准备好在 SageMaker 上运行数据并行训练。

在以下代码块中,我们定义了 TensorFlow 版本(2.8),Python 版本(3.9),实例类型以及实例的数量。此外,我们还传递了若干训练超参数。由于 MNIST 数据集已经作为训练脚本的一部分从互联网下载,因此无需将数据传递给 estimator_ms.fit() 方法:

from sagemaker.tensorflow import TensorFlow
ps_instance_type = 'ml.p2.xlarge'
ps_instance_count = 2
hyperparameters = {'epochs': 4, 'batch-size-per-device' : 16, 'steps-per-epoch': 100}
estimator_ms = TensorFlow(
                       source_dir='1_sources',
                       entry_point='train_ms.py', 
                       role=role,
                       framework_version='2.8',
                       py_version='py39',
                       disable_profiler=True,
                       debugger_hook_config=False,
                       hyperparameters=hyperparameters,
                       instance_count=ps_instance_count, 
                       instance_type=ps_instance_type,
                       )
estimator_ms.fit()

使用默认设置,训练任务应该在 10 到 12 分钟内完成。你可以随意尝试集群中的节点数量和实例类型,并观察 'TF_CONFIG'、训练速度和收敛性的变化。

在接下来的部分中,我们将学习一种数据并行的开源替代方案——Horovod 框架。

使用 Horovod 框架

Horovod 框架为最流行的深度学习框架提供了同步数据并行的实现,如 TensorFlow 1 和 TensorFlow 2(包括 Keras)、PyTorch 和 Apache MXNet。Horovod 的一个优点是,它只需要最少的修改即可分配训练任务,这兼容各种集群布局。Horovod 支持几种通信后端:Gloo 和 Open MPI 用于基于 CPU 的训练,NCCL 用于运行在 NVIDIA GPU 设备上的训练。

Horovod 配备了一些功能,旨在解决我们之前讨论的 Allreduce 算法的概念限制。为了减少 allreduce 计算过程中的等待时间,并提高训练设备的利用率,Horovod 引入了一个名为 allreduceallgather 的概念,并将其组织成层次结构,从而实现更好的整体性能。此外,Horovod 还提供了一个 Autotune 工具,可以通过调整训练参数来调优训练作业的性能。请注意,运行 Autotune 作业并不适合用于生产环境。

现在,让我们回顾一下如何在 SageMaker 上使用 Horovod 进行 TensorFlow 2 的训练。请注意,Horovod 原生支持 TensorFlow 和 PyTorch 框架。在本章中,我们将只回顾 TensorFlow 2 的 Horovod 实现,因为 PyTorch 的实现将非常相似。我们将解决之前提到的 MNIST 分类问题。

配置 Horovod 集群

与 MWMS 不同,我们不需要在训练脚本中配置和设置训练集群,因为 Horovod 已被 SageMaker 支持。Horovod 集群配置是在 TensorFlow Estimator API 级别通过distribution对象完成的,如以下代码块所示:

distribution = {"mpi": {"enabled": True, "custom_mpi_options": "-verbose --NCCL_DEBUG=INFO", "processes_per_host": 1}}

注意processes_per_host参数,它应该与所选实例类型中的 GPU 数量相匹配。此外,你可以根据需要设置custom_mpi_options,这些选项会传递给mpirun运行工具。你可以在www.open-mpi.org/doc/v4.0/man1/mpirun.1.php查看支持的 MPI 选项列表。

开发训练脚本

你可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter6/1_sources/train_hvd.py找到完整的训练脚本。让我们执行以下步骤:

  1. 我们通过_initiate_hvd()方法在训练脚本中启动 Horovod。我们还需要将 Horovod 训练进程与可用的 GPU 设备关联起来(每个进程一个设备):

    def _initiate_hvd():
        hvd.init()
        gpus = tf.config.experimental.list_physical_devices("GPU")
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        if gpus:
    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], "GPU")
    
  2. 接下来,我们需要根据世界大小对数据集进行切片,以便每个进程可以根据其全局排名获取数据切片。为此,我们使用 TensorFlow 数据集实例的shard方法。请注意,我们正在通过 Horovod 的size()rank()属性获取给定训练进程的本地和全局排名:

    train_dataset = train_dataset.shard(hvd.size(), hvd.rank())
    
  3. 然后,我们使用DistributedOptimizer Horovod 包装器来启用分布式梯度更新。请注意,我们包装的是本地 TF2 优化器的一个实例:

    optimizer = tf.keras.optimizers.SGD(learning_rate=0.001 * hvd.size())
    optimizer = hvd.DistributedOptimizer(optimizer)
    
  4. 最后,我们使用特殊的 Horovod 回调函数,Keras 将在训练循环中使用这些回调函数:

    • 使用hvd.callbacks.BroadcastGlobalVariablesCallback(0)将来自rank=0进程的初始变量分发到集群中的其他训练进程。

    • 使用hvd.callbacks.MetricAverageCallback()来计算所有训练进程的全局平均指标。

  5. 这些回调函数随后将传递给model.fit()方法,如下所示:

        hvd_model.fit(
            shareded_by_rank_dataset,
            epochs=args.epochs,
            steps_per_epoch=args.steps_per_epoch // hvd.size(),
            callbacks=callbacks,
        )
    

这些是你训练脚本中最少的改动,它们可以让你使用 Horovod。

运行 SageMaker 作业

SageMaker 训练作业配置与 MWMS 示例类似,但我们将添加distribution参数,允许我们设置 MPI 参数并定义每个主机启动多少个进程:

from sagemaker.tensorflow import TensorFlow
ps_instance_type = 'ml.p2.xlarge'
ps_instance_count = 2
distribution = {"mpi": {"enabled": True, "custom_mpi_options": "-verbose --NCCL_DEBUG=INFO", "processes_per_host": 1}}
hyperparameters = {'epochs': 4, 'batch-size-per-device' : 16, 'steps-per-epoch': 100}
estimator_hvd = TensorFlow(
                       source_dir='1_sources',
                       entry_point='train_hvd.py', 
                       role=role,
                       framework_version='2.8',
                       py_version='py39',
                       disable_profiler=True,
                       debugger_hook_config=False,
                       hyperparameters=hyperparameters,
                       instance_count=ps_instance_count, 
                       instance_type=ps_instance_type,
                       distribution=distribution
                       )
estimator_hvd.fit()

在这里,我们实现了使用 TensorFlow 2 MWMS 和 TensorFlow 2 Horovod 的数据并行训练任务的最小可行示例。现在,你应该已经具备了开发基准训练任务的一些实践经验。两种 Allreduce 实现中都有更多的调节选项和功能,我们鼓励你在实际使用案例中探索和尝试。选择具体的实现方式(MWMS 或 Horovod)在许多情况下是基于使用场景的,没有明确的优劣之分。Horovod 的优势在于它支持多个深度学习框架,并且其成熟性(特别是在故障排除和优化工具方面)。另一方面,TensorFlow 2 策略与各种 TensorFlow API 和不同的方法原生集成,其中许多目前处于实验模式。

在接下来的章节中,我们将转向 PyTorch 框架,并回顾其原生数据并行实现。

工程化 PyTorch 数据并行训练

PyTorch 提供了一个原生实现的数据并行工具 torch.distributed.run,用于简化和协调进程启动。与 Horovod 类似,PyTorch DDP 支持 NCCL、Gloo 和 MPI 通信后端。此外,PyTorch DDP 原生支持 混合精度自动混合精度 (AMP),这允许你以半精度训练模型,并对模型精度和训练收敛性的影响最小化。AMP 的好处包括加速计算并减少内存占用。

尽管 SageMaker 本身不原生支持 PyTorch DDP,但仍然可以在 SageMaker 上运行 DDP 训练任务。让我们回顾一下实现示例。

我们采用预训练的 CV Resnet18 模型,并对其进行微调,以便对蚂蚁和蜜蜂进行分类。我们使用数据并行来将任务分配到两台 p2.xlarge 实例中,每个实例有一个 GPU 设备。可以随意更改或修改训练集群中实例的数量和类型,并观察这如何改变训练速度。

请注意,这只是小规模训练,不能代表实际任务中训练效率的表现。

接下来,我们将重点介绍关键代码构建块。一个笔记本和其他代码资源可以在 github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter6/2_distributed_training_PyTorch.ipynb 获取。

启动训练进程

Amazon SageMaker 原生不支持 PyTorch DDP 训练。具体来说,它不知道如何在训练集群中启动分布式 DDP 进程。因此,我们需要开发一个启动工具来执行此功能。这个工具非常简单,并且可以在任何其他基于 DDP 的训练任务中复用。

在启动脚本中,我们将使用 DDP 模块 torch.distributed.run,它简化了在集群中生成训练进程的过程。作为启动脚本的一部分,我们需要收集训练世界的信息,特别是集群中的节点数量和 GPU 设备数量,并确定哪个节点将作为主协调者。然后,torch.distributed.run 会生成多个训练进程。请参考图 6.5以获得可视化说明:

图 6.5 – 在 N 个节点上启动 PyTorch DDP 训练,使用两个 GPU

图 6.5 – 在 N 个节点上启动 PyTorch DDP 训练,使用两个 GPU

让我们突出几个启动脚本中的关键部分:

  1. 首先,我们需要收集有关 SageMaker 训练集群的信息。为此,我们使用 SageMaker 自动设置的环境变量:

        nodes = json.loads(os.getenv("SM_HOSTS"))
        nnodes = len(nodes)
        node_rank = nodes.index(os.getenv("SM_CURRENT_HOST"))
        nproc_per_node = os.getenv("SM_NUM_GPUS", 1)
    
  2. 接下来,我们需要构建命令行以启动 torch.distributed.run

        cmd = [
            sys.executable,
            "-m",
            "torch.distributed.run",
            f"--nproc_per_node={nproc_per_node}",
            f"--nnodes={str(nnodes)}",
            f"--node_rank={node_rank}",
            f"--rdzv_id={os.getenv('SAGEMAKER_JOB_NAME')}",
            "--rdzv_backend=c10d",
            f"--rdzv_endpoint={nodes[0]}:{RDZV_PORT}",
            distr_args.train_script,
        ]
        # Adding training hyperparameters which will then be passed in training script
        cmd.extend(training_hyperparameters)
    

请注意,我们在命令行末尾添加了“原样”的 training hyperparameters。这些参数不是由启动器处理的,而是由训练脚本来配置训练。

  1. 最后,我们使用 Python 的 subprocess.Popen 启动 torch.distributed.run 工具作为模块:

        process = subprocess.Popen(cmd, env=os.environ)
        process.wait()
        if process.returncode != 0:
            raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
    

请注意,我们将环境变量复制到子进程,以保留所有 SageMaker 变量。如果生成的进程返回非零代码(表示错误),我们将引发异常,将错误代码传播到 SageMaker 控制平面。

总结来说,我们的启动工具负责收集训练集群配置,然后在每个节点上启动 torch.distributed.run。该工具随后负责在每个节点启动多个训练进程。

为 DDP 采用训练脚本

要使用 DDP,我们需要对训练脚本进行最小的修改:

  1. 首先,我们初始化训练进程并将其添加到 DDP 进程组:

    dist.init_process_group(
        backend="nccl",
        rank=int(os.getenv("RANK", 0)),
        world_size=int(os.getenv("WORLD_SIZE", 1)),
    )
    

由于我们使用基于 GPU 的实例,我们使用 NCCL 通信后端。此外,我们利用 torch.distributed.run 模块设置的环境变量:世界大小和全局排名。

  1. 接下来,我们需要确定哪个 GPU 设备将存储模型并进行计算。我们使用 torch.distributed.run 在进程生成时设置的 LOCAL_RANK 变量:

    torch.cuda.set_device(os.getenv("LOCAL_RANK"))
    device = torch.device("cuda")
    model = model.to(device)
    
  2. 然后,我们将常规的 PyTorch 模型包装成一个特殊的 DDP 实现。这个实现允许我们像处理常规本地存储的模型一样使用 PyTorch 模型。在背后,DDP 模块实现了在进程组中各训练进程之间的梯度同步。同时,请注意,我们正在根据世界大小缩小用户提供的全局批量大小:

    model = DDP(model)
    args.batch_size //= dist.get_world_size()
    args.batch_size = max(args.batch_size, 1)
    
  3. 我们需要做的最后一步是修改训练数据加载器,以确保每个训练进程在训练步骤中获取唯一的数据切片。为此,我们使用 DistributedSampler,它根据进程总数采样数据记录,并处理全局排名:

        train_sampler = torch.utils.data.distributed.DistributedSampler(
            image_datasets["train"], num_replicas=args.world_size, rank=args.rank
        )
        train_loader = torch.utils.data.DataLoader(
            image_datasets["train"],
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=0,
            pin_memory=True,
            sampler=train_sampler,
        ) 
    

训练脚本的其余部分与非分布式训练类似。如你所见,为使训练脚本兼容 PyTorch DDP 所做的修改非常少。

运行 SageMaker 训练任务

一旦启动器和训练脚本准备好,我们就可以开始 SageMaker 训练任务。请注意,我们将启动器脚本指定为entry_point参数。训练脚本的引用与训练超参数一起提供,放在hyperparameter对象中:

from sagemaker.pytorch import PyTorch
ps_instance_type = 'ml.p3.2xlarge'
ps_instance_count = 2
hyperparameters = {
  'train-script': 'train_ddp.py',
  'epochs': 25,
  }
estimator_ms = PyTorch(
                       source_dir='2_sources',
                       entry_point='launcher.py', 
                       role=role,
                       framework_version='1.9',
                       py_version='py38',
                       disable_profiler=True,
                       debugger_hook_config=False,
                       hyperparameters=hyperparameters,
                       instance_count=ps_instance_count, 
                       instance_type=ps_instance_type,
                       )
estimator_ms.fit(inputs={"train":f"{data_url}/train", "val":f"{data_url}/val"})

训练任务应在 8 到 9 分钟内完成。你可以随时查看训练任务日志中的调试信息。此外,你还可以尝试其他参数,例如实例类型和大小、训练轮次、批量大小等。

在本节中,我们学习了如何在 PyTorch 框架中使用原生的数据并行实现。在下一节中,我们将介绍 SageMaker 的专有数据并行实现。

工程化 SageMaker 的 DDP 任务

SDDP 库提供了一个专有的数据并行实现,并与其他 SageMaker 功能原生集成。SDDP 被打包在 SageMaker DL 容器中,并支持 TensorFlow 2 和 PyTorch 框架。

SDDP 利用 MPI(类似于 Horovod)来管理训练集群中的进程。在背后,SDDP 使用ml.p3.16xlargeml.p3dn.24xlargeml.p4d.24xlarge。SDDP 提供了一个与 Horovod 和 PyTorch DDP 非常相似的 API,这使得从开源实现切换到它变得非常容易。

SDDP 实现了一个修改版的 Allreduce 算法,并进行了多项优化,以提高整体训练性能,特别是在allreduce操作中的等待时间。如前所述,在同步 Allreduce 算法中,通常分布式的allreduce操作是瓶颈,并且随着训练集群规模的扩大,它变得更加低效。请查看图 6.6

图 6.6 – 集群增加时的 Allreduce 时间

图 6.6 – 集群增加时的 Allreduce 时间

为了提高训练效率,特别是在大规模集群中,SDDP 引入了几项新的优化:

  • SDDP 在训练过程中利用 GPU 和 CPU 设备,因此 GPU 设备执行前向和反向传播,CPU 设备则在allreduce阶段执行梯度平均和与其他训练进程的通信。这种方法使得计算操作和allreduce能够并行运行,从而最大化资源利用率。

  • SDDP 支持allreduce(例如 Horovod 的张量融合功能)。

因此,AWS 声称 SDDP 随着训练集群规模的增加,训练吞吐量几乎实现线性扩展。AWS 发布了以下基准测试,展示了 SDDP 与原生 PyTorch DDP 相比的优化收益:对于 8 节点的p3dn.24xl集群,在训练 BERT 模型时,SDDP 比 PyTorch DDP 快 41%,在训练 MaskRCNN 模型时快 13%。更多细节请参考此文献:docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-intro.xhtml

在构建 SDDP 训练任务时,请牢记以下几个方面:

  • SDDP 依赖于 CPU 设备来执行allreduce操作。大多数框架的数据加载器也使用 CPU。因此,请确保控制你的 CPU 使用,以避免过度利用。在第七章《深度学习训练的操作化》中,我们将讨论可以用来控制资源利用的工具,例如 SageMaker 调试器。或者,你可以将数据加载操作转移到 GPU 上。然而,在这种情况下,你将有较少的 GPU 内存来加载模型并执行其前向和反向传递。

  • 在小型集群或单个节点上使用 SDDP 可能没有显著的效果或任何好处,因为它的设计目标是专门解决大规模训练集群的瓶颈。

让我们回顾一下基于 SDDP 的训练任务示例。为此,我们将重用之前的 PyTorch DDP,并做最小的修改,从 PyTorch DDP 切换到 SDDP 库。

作为一个训练任务,我们使用与 PyTorch DDP 示例中相同的二分类 CV。由于 SDDP 本身得到了 SageMaker 的原生支持,我们无需开发任何自定义启动工具。SDDP 使用mpirun工具在集群中启动训练进程。你可以使用distribution参数启用数据并行执行,并提供任何mpi选项,如下所示:

distribution = { 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}

现在,让我们开始采用训练脚本。

采用训练脚本

SDDP 的起始版本 1.4.0 是一个集成的 PyTorch DDP 包,我们在前面的示例中将其作为特定的后端选项使用。这样可以显著减少使用 SDDP 所需的更改。事实上,如果你已经有一个启用了 DDP 的训练脚本,你只需要添加torch_sddp包的导入,并在初始化进程组时使用smddp通信后端,如下所示:

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

请注意,SDDP v1.4 仅在最新的 PyTorch v10 DL 容器中可用。对于早期版本,SDDP API 略有不同。更多细节请参考官方 API 文档:sagemaker.readthedocs.io/en/stable/api/training/distributed.xhtml#the-sagemaker-distributed-data-parallel-library

运行 SDDP SageMaker 训练任务

启动 SDDP 任务时,需要提供一个特殊的distribution对象,并配置数据并行性。另一个需要注意的事项是,SDDP 仅适用于有限的多 GPU 实例类型:ml.p3.16xlargeml.p3dn.24xlargeml.p4d.24xlarge。请参考以下内容:

from sagemaker.pytorch import PyTorch
instance_type = 'ml.p3.16xlarge'
instance_count = 2
distribution = { 
    "smdistributed": { 
        "dataparallel": {
            "enabled": True, 
            "custom_mpi_options": "-verbose -x NCCL_DEBUG=VERSION"
        }
    }
}
sm_dp_estimator = PyTorch(
          entry_point="train_sm_dp.py",
          source_dir='3_sources',
          role=role,
          instance_type=instance_type,
          sagemaker_session=sagemaker_session,
          framework_version='1.10',
          py_version='py38',
          instance_count=2,
          hyperparameters={
              "batch-size":64,
              "epochs":25,
          },
          disable_profiler=True,
          debugger_hook_config=False,
          distribution=distribution,
          base_job_name="SM-DP",
      )

请注意,由于我们使用的是一个小型数据集,这个训练样本无法体现与开源数据并行框架相比,SDDP 的任何性能优势。

数据并行性的总结

到目前为止,我们讨论了如何加速训练那些能够适配单一设备内存的深度学习模型。我们在讨论和开发训练脚本时,使用了深度学习框架的本地实现、开源解决方案以及专有的跨框架 Allreduce 实现(分别是 Horovod 和 SageMaker SDDP)。然而,我们并未对给定实现的训练效率进行基准测试。虽然每个使用场景都是独特的,但通常的建议是,当涉及大规模且时间较长的训练过程,尤其是需要大规模集群时,应优先考虑使用 SDDP。如果你有中型或小型的训练任务,仍然可以考虑使用框架本地的数据并行实现。在这种情况下,SDDP 的性能优势可能微乎其微。

在接下来的章节中,我们将讨论如何使用模型并行优化训练那些无法完全适配单个 GPU 内存的模型。

工程化模型并行训练任务

在模型并行中,模型的单一副本会被分布到两个或更多的训练设备上,以避免单个 GPU 设备的内存限制。模型并行的简单方法是将模型的各层显式地分配到不同的设备上。在这种情况下,前向计算将在存储第一组层的 GPU 设备上进行。然后,结果将被传输到存储下一组层的 GPU 设备上,依此类推。在反向传递过程中,层之间的交接会按相反顺序发生。这种类型的模型并行被称为简单模型并行垂直模型并行,因为我们将模型在设备之间进行垂直切分。然而,这种模型并行方式效率较低,因为每个 GPU 设备都需要等待其他设备完成计算,导致显著的时间浪费。更高效的模型并行方式叫做流水线并行。这种方法将单一的数据批次分成多个微批次,并通过重叠不同微批次的梯度计算,尽量减少等待时间。请参见图 6.7,对比简单模型并行和流水线模型并行:

图 6.7 – 简单模型并行和流水线模型并行

图 6.7 – 简单模型并行和流水线模型并行

图表来源

ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.xhtml

实现流水线并行面临一些挑战,因为你可能需要重新实现训练脚本,将模型的各个部分分配到不同的设备上,并在训练循环中反映新的计算流程。此外,你还需要决定如何在同一节点内和跨节点之间优化地放置模型层。还有,流水线并行不支持条件流,并要求每一层接受一个张量作为输入并产生一个张量作为输出。你还需要为每种新的模型架构重新实现流水线并行。接下来在本节中,我们将看到 SMDP 库如何解决这些挑战。

垂直拆分模型是一种最小化内存占用的方法。另一种并行化方法是称为张量并行。每个张量(数据输入和层输出)会被拆分到多个设备上并行处理。然后,将单独的结果进行聚合。张量并行是可行的,因为许多计算操作可以表示为矩阵运算,而矩阵运算可以沿着X轴或Y轴进行拆分。请参见图 6.8,它提供了张量如何拆分的可视化表示。张量并行也被称为水平并行

图 6.8 – 行式和列式张量并行

图像来源

github.com/huggingface/transformers/blob/main/docs/source/en/perf_train_gpu_many.mdx

流水线并行和张量模型并行可以结合使用。此外,还可以加入数据并行,以实现更高的并行化和更好的训练速度。数据并行与模型并行的结合被称为混合并行。这种方法被用来训练当前大多数最先进的 SOTA NLP 模型,如 T5 或 GPT3。请参见图 6.9,它展示了流水线并行和数据并行的结合:

图 6.9 – 结合流水线并行和数据并行

图 6.9 – 结合流水线并行和数据并行

现在我们已经刷新了对关键模型并行方法的理解,接下来让我们回顾一下 SDMP 库——SageMaker 专有的模型并行实现。

使用 SDMP 进行工程训练

SDMP 是一个功能丰富的库,能够实现各种类型的模型并行性和混合并行性,并针对 SageMaker 基础设施进行了优化。它支持 TensorFlow 和 PyTorch 框架,并允许你在最少的代码更改下,自动将模型划分到不同的设备。像 SDDP 一样,SDMP 使用 MPI 在训练集群中协调任务,在 GPU 设备上进行前向和反向计算,在 CPU 设备上进行通信任务。

SDMP 具有一些显著特点,可以简化模型并行训练作业的开发,并优化训练时硬件的利用率:

  • SDMP 支持任意的模型架构,并且对训练脚本的修改最小,不会产生任何精度上的惩罚。

  • 自动化模型拆分将在训练集群中将模型划分到不同设备之间。你可以选择优化速度和内存利用率。此外,SDMP 还支持手动模型拆分(然而,在实践中,这种方法通常不是一个好的选择)。

  • 交错管道是对简单模型管道的改进,允许你通过优先执行反向操作来最小化处理微批次的时间:

图 6.10 – 简单和交错管道的比较

图 6.10 - 简单和交错管道的比较

图的来源

docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features.xhtml

虽然 SDMP 官方支持 TensorFlow 2 和 PyTorch,但某些优化功能仅对 PyTorch 可用。针对 PyTorch 的扩展支持包括以下内容:

  • 优化器状态分片允许你不仅将模型分割,还可以在数据并行组的训练设备之间分割优化器状态。这进一步减少了每个设备在训练过程中的内存占用。请注意,优化器状态分片会增加一个聚合步骤(当全局优化器状态从各个分片重新构建时),这会导致额外的延迟。

  • 除了管道和数据并行性之外,张量并行性也非常重要。张量并行性对于那些无法完全适配到单一 GPU 设备的大层次(如嵌入层)尤为有效。

  • 激活检查点激活卸载是另外两种技术,它们通过增加一些计算时间来重建训练状态,从而进一步减少训练内存占用。

由于使用这些高级优化功能需要在内存和计算之间进行权衡,通常建议仅在大型模型(即数十亿参数)上使用它们。

现在,让我们使用 SDMP 库开发一个混合并行作业。我们将重用之前的 PyTorch 示例和 CV 模型。

注意

这个例子仅具有教学目的。通常,CV 模型(如 Resnet18)可以适应单个 GPU,在这种情况下,不需要模型并行。然而,对于演示目的,较小的模型更易于管理并且训练速度较快。

配置模型和混合并行性

首先,让我们了解一下我们的训练将如何执行以及如何配置并行性。为此,我们将使用 SageMaker 训练作业的分发对象。它有两个关键组件:model parallelmpi

SageMaker 依赖 mpi 工具来运行分布式计算。在以下代码片段中,我们设置它运行 8 个训练进程。这里,processes_per_host 定义了每个主机将运行多少个训练进程,这些进程包括运行模型并行、数据并行或张量并行的进程。在大多数情况下,进程数量应与节点中可用的 GPU 数量相匹配。

Modelparallel 对象定义了 SDMP 库的配置。然后,在代码片段中,我们设置了 2 路模型并行(partitions 参数设置为 2)。同时,我们通过将 ddp 参数设置为 True 来启用数据并行。当启用数据并行时,SDMP 将根据训练进程数和模型并行大小自动推断数据并行大小。另一个重要的参数是 auto_partition,因此 SDMP 会自动在 GPU 设备之间划分模型。

在以下代码块中,我们配置了我们的训练作业,使其在 2 个实例上运行,总共有 16 个 GPU。我们的 distribution 对象定义了 2 路模型并行。由于总训练进程数为 16,SDMP 将自动推断出 8 路数据并行。换句话说,我们将模型划分为 2 个 GPU 设备,并且总共有 8 个模型副本:

smd_mp_estimator = PyTorch(
# ... other job parameters are reducted for brevity
instance_count=2,
instance_type= 'ml.p3.16xlarge',          
distribution={
                  "modelparallel": {
                      "enabled":True,
                      "parameters": {
                          "microbatches": 8, 
                          "placement_strategy": "cluster", 
                          "pipeline": "interleaved",
                          "optimize": "speed", 
                          "partitions": 2,
                          "auto_partition": True,
                          "ddp": True,
                      }
                  }
              },
            "mpi": {
                    "enabled": True,
                    "processes_per_host": 8,
                    "custom_mpi_options": mpioptions 
              }

请注意,您需要将混合并行的配置与集群布局(节点数和 GPU 设备数)对齐。SageMaker Python SDK 提供了混合并行配置的预验证;然而,这并不能保证在训练过程中将使用所有 GPU 设备。最好在训练脚本中添加调试信息,以确保所有 GPU 设备都得到正确利用。

采用训练脚本

SDMP 的一个好处是它对训练脚本的改动最小。这是通过使用 Python 装饰器来定义需要以模型并行或混合方式运行的计算实现的。此外,SDMP 提供了类似 Horovod 或 PyTorch DDP 等其他分布式库的 API。在以下代码块中,我们仅突出显示了关键部分。完整的源代码可以在github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/tree/main/chapter6/4_sources找到:

  1. 我们从导入和初始化 SDMP 库开始:

    import smdistributed.modelparallel.torch as smp
    smp.init()
    
  2. 一旦库初始化完成,我们可以使用 SDMP API 来检查我们的混合并行是否已正确配置。为此,您可以在训练脚本中运行以下debug语句:

    logger.debug(
    f"Hello from global rank {smp.rank()}. "
          f"Local rank {smp.local_rank()} and local size {smp.local_size()}. "
          f"List of ranks where current model is stored {smp.get_mp_group()}. "
          f"List of ranks with different replicas of the same model {smp.get_dp_group()}. "
          f"Current MP rank {smp.mp_rank()} and MP size is {smp.mp_size()}. "
            f"Current DP rank {smp.dp_rank()} and DP size is {smp.dp_size()}."
        )
    
  3. 输出将在每个训练过程中生成。让我们查看来自全局排名0的输出。在这里,方括号中的消息前缀由 MPP 工具提供,标记了唯一的 MPI 进程,algo-1是主机名的引用。从调试信息中,您可以确认我们已经配置了 2 路并行和 8 路数据并行。此外,我们还可以观察到数据并行和模型并行组的 GPU 分配情况:

    [1,mpirank:0,algo-1]:INFO:__main__:Hello from global rank 0\. Local rank 0 and local size 8\. List of ranks where current model is stored [0, 1]. List of ranks with different replicas of the same model [0, 2, 4, 6, 8, 10, 12, 14]. Current MP rank 0 and MP size is 2\. Current DP rank 0 and DP size is 8.
    
  4. SDMP 管理模型分区到 GPU 设备的分配,您不需要显式地将模型移动到特定设备(在常规的 PyTorch 脚本中,您需要通过调用model.to(device)方法显式移动模型)。在每个训练脚本中,您需要根据 SMDP 本地排名选择一个 GPU 设备:

    torch.cuda.set_device(smp.local_rank())
    device = torch.device("cuda")
    
  5. 接下来,我们需要将 PyTorch 模型和优化器包装在 SDMP 实现中。这是为了在模型并行和数据并行组之间建立通信。

  6. 一旦包装完成,您将需要在训练脚本中使用 SDMP 包装的模型和优化器版本。请注意,您仍然需要使用 PyTorch 的input_tensor.to(device)方法将输入张量(例如,数据记录和标签)移动到此设备:

    model = smp.DistributedModel(model)
    optimizer = smp.DistributedOptimizer(optimizer)
    
  7. 之后,我们需要配置我们的数据加载器。SDMP 对数据加载器没有特定要求,除了确保批次大小一致。建议您使用drop_last=True标志来强制执行这一点。因为在内部,SDMP 将批量数据分解为一组微批次来实现流水线处理。因此,我们需要确保批次大小始终能被微批次大小整除。请注意,在下面的代码块中,我们使用 SDMP API 配置了一个用于数据并行的分布式采样器:

        dataloaders_dict = {}
        train_sampler = torch.utils.data.distributed.DistributedSampler(
            image_datasets["train"], num_replicas=sdmp_args.dp_size, rank=sdmp_args.dp_rank)
        dataloaders_dict["train"] = torch.utils.data.DataLoader(
            image_datasets["train"],
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=0,
            pin_memory=True,
            sampler=train_sampler,
            drop_last=True,
        )
        dataloaders_dict["val"] = torch.utils.data.DataLoader(
            image_datasets["val"],
            batch_size=args.batch_size,
            shuffle=False,
            drop_last=True,
        )
    
  8. 一旦我们配置了模型、优化器和数据加载器,就可以编写训练和验证循环了。为了实现模型并行,SDMP 提供了一个@smp.step装饰器。任何使用@smp.step装饰的函数都会以流水线方式执行内部计算。换句话说,它将批量数据拆分为一组微批次,并协调在 GPU 设备之间的模型分区的计算。在这里,训练和测试计算都使用了@smp.step装饰器。请注意,训练步骤包含了前向和反向传播,因此 SDMP 可以在所有分区上计算梯度。在测试步骤中只有前向传播:

    @smp.step
    def train_step(model, data, target, criterion):
        output = model(data)
        loss = criterion(output, target)
        model.backward(loss)  #  instead of PyTorch loss.backward()
        return output, loss
    @smp.step
    def test_step(model, data, target, criterion):
        output = model(data)
        loss = criterion(output, target)
        return output, loss
    

注意另一个区别:在计算损失时,我们使用了model.backward(loss)的 SDMP 方法。因此,SDMP 可以正确计算跨模型分区的梯度值。

  1. 我们在外部训练循环中使用了装饰过的训练和测试步骤,如下所示。训练循环的构造类似于典型的 PyTorch 训练循环,唯一的区别是,由于 SDMP 对微批次实现了流水线操作,损失值也会对微批次进行计算(即loss_mb变量)。因此,为了计算整个批次的平均损失,我们调用reduce_mean()方法。注意,所有由@smp.step装饰的函数返回的变量都是提供便捷 API 的类实例,可以跨小批次进行操作(例如.reduce_mean().concat()方法):

    for epoch in range(num_epochs):
            for phase in ["train", "val"]:
                if phase == "train":
                    model.train()  # Set model to training mode
                else:
                    model.eval()  # Set model to evaluate mode
                for inputs, labels in dataloaders[phase]:
                    inputs = inputs.to(device)
                    labels = labels.to(device)
                    optimizer.zero_grad()
                    with torch.set_grad_enabled(phase == "train"):
                        if phase == "train":
                            outputs, loss_mb = train_step(model, inputs, labels, criterion)
                            loss = loss_mb.reduce_mean()
                            optimizer.step()
                        else:
                            outputs, loss_mb = test_step(model, inputs, labels, criterion)
                            loss = loss_mb.reduce_mean()
    
  2. 训练完成后,我们需要保存分布式模型。为此,SMDP 提供了smp.save()方法,支持以 pickle 格式保存模型和优化器的状态。你可以通过使用partial标志来选择是否持久化模型分区。如果启用了部分保存,则模型分区将与其模型并行排名一起单独保存。在以下代码块中,我们保存了一个单一的模型检查点。请注意,我们基于排名过滤器在单个进程中保存模型,以避免冲突:

        if smp.dp_rank() == 0:
            model_file_path = os.path.join(
                os.environ["SM_MODEL_DIR"], f"finetuned-{args.model_name}-checkpoint.pt"
            )
            model_dict = model.state_dict()  # save the full model
            opt_dict = optimizer.state_dict()  # save the full optimizer state
            smp.save(
                {"model_state_dict": model_dict, "optimizer_state_dict": opt_dict},
                model_file_path,
                partial=False,
            )
    
  3. 一旦我们的测试完成,SageMaker 将把模型和优化器的检查点上传到 S3 位置。你可以按如下方式使用此模型进行推理:

    model_state = torch.load('finetuned-resnet-checkpoint.pt')['model_state_dict']
    model_ft = models.resnet18(pretrained=False)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, num_classes)
    model_ft.load_state_dict(model_state)
    outputs = model_ft(inputs)
    

这些是训练脚本中所需的最小更改,以使其与 SDMP 库兼容,从而实现模型并行或混合并行。我们鼓励你尝试各种 SDMP 配置参数(请参考前一部分中的distribution对象),以便培养良好的直觉,具体包括以下内容:

  • 改变模型的分区数量。我们的实现有 2 个分区;你可以尝试设置 1 或 4 个分区,看看这如何改变数据并行和模型并行的分组。

  • 改变微批次批次大小的数量,看看它如何影响训练速度。在生产场景中,你可能需要探索批次大小和微批次的上限,以提高训练效率。

  • 查看管道实现的类型——交错式或简单式——如何影响训练速度。

额外的考虑事项

我们使用了相对简单和小型的模型,如 Resnet,来演示如何实现混合并行。然而,像 GPT-n 这样的更复杂模型的实现将需要额外的考虑。以下部分将详细说明这些内容。

张量并行

张量并行仅在 SDMP 库的 PyTorch 版本中可用。张量并行适用于那些参数层消耗大量 GPU 内存的场景(例如嵌入表)。使用张量并行时,您需要确保 SDMP 支持您的模型的模块。SDMP 提供了常见模块的分布式实现,如 nn.Linearnn.Embedding 等。如果某个特定模块不被支持,首先,您需要实现一个张量并行版本。有关详细信息,请参考 SDMP API 文档:sagemaker.readthedocs.io/en/stable/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.xhtml

参考实现

AWS 提供了一些示例脚本,用于使用 SDMP 库训练流行的大型模型,如 GPT2、GPT-J 和 BERT。请查看官方 GitHub 仓库:github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel

您还可以在 docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-best-practices.xhtml 找到 SDMP 的参考配置。

到目前为止,我们已经涵盖了分配训练作业的不同方式,并查看了示例实现。在接下来的部分,我们将介绍一些可以提高分布式训练作业效率的参数。

优化分布式训练作业

在许多情况下,优化大规模分布式训练作业需要大量的试验和错误,并且非常依赖于运行时环境、硬件栈、模型架构以及其他参数。但有一些关键的调整项可以帮助优化您的训练速度和整体效率。

集群布局与计算亲和性

在运行分布式训练作业时,特别是模型并行或混合并行时,理解不同类型的计算如何与您的集群布局对齐始终是一个好主意。

假设我们需要在一个包含两个节点、每个节点有两个 GPU 设备的集群中运行模型并行训练。总共的训练进程数为 4,全球排名从 0 到 3,局部排名为 0 和 1。我们假设模型副本和梯度可以适配到两个设备中。在这种情况下,我们需要确保每个模型会被存储在每个节点内:一个模型副本会放置在排名 0 和 1(局部排名 0 和 1)上,另一个放置在排名 2 和 3(局部排名 0 和 1)上。这将确保模型各层之间的通信通过更快的 GPU 互联连接进行,并且通常不会穿越较慢的节点内网络。

为了解决这个问题,SDMP 提供了一个特殊的参数 placement_strategy,允许你控制训练过程与硬件的亲和性。

通信后端

在本章中,我们介绍了一些最流行的通信后端,例如 NCCL、Gloo 和 MPI。以下列表是根据具体情况选择后端时的一个经验法则:

  • 消息传递接口MPI)是分布式计算中的一种通信标准,具有多个后端实现,例如 Open-MPI、MVAPICH2 等。MPI 后端还支持在 CUDA 张量上进行 GPU 之间的操作。然而,如果你有其他选择,MPI 通常不是训练任务的最佳选择。

  • Gloo 后端广泛支持 CPU 设备之间的点对点和集体计算,以及 GPU 设备之间的集体计算。Gloo 是在 CPU 设备上进行初步调试的不错选择。然而,通常在使用 GPU 设备进行训练时,你应优先选择 NCCL。

  • NCCL 后端由 NVIDIA 提供,最适合在 NVIDIA GPU 上进行训练任务。

  • 自定义 后端可以作为深度学习框架的一部分提供。例如,TensorFlow 2 提供了一个自定义的 RING 后端。

注意

在使用较新的 SOTA 模型时,确保你选择的通信后端支持模型架构所需的集体操作和点对点操作。

训练超参数

有许多训练超参数会影响训练效率。虽然我们不会涵盖所有的超参数,但我们列出了一些你可以在优化过程中调整的超参数:

  • 使用 AMP 来减少内存需求,并在对准确性和训练收敛性影响最小的情况下加速训练。AMP 是一种流行的技术,它在前向、反向和更新步骤中结合了单精度(FP32)和半精度(FP16)张量。请注意,为了在使用 AMP 时获得显著的改进,通常需要使用较大的批量大小。

  • 使用 硬件优化的数据类型(如 TF32、BF32 和 BF16)来加速训练。这些数据类型针对特定的深度学习计算进行了优化,相较于常见的 FP32 和 FP16 类型提供了加速效果。请注意,要使用这些类型,需确保你的框架、模型架构和硬件都支持它们。

  • 优化你的全局批次大小以加速训练。随着训练集群的扩展,确保相应地更新全局批次大小。通常,局部批次大小的上限由可用的 GPU 内存定义(如果局部批次大小无法适应内存,你可能会看到CUDA OOM错误)。请记住,将批次大小增加到某个阈值以上,可能不会提高全局训练吞吐量。你可以在NVIDIA指南中找到一些额外的资料和基准测试:docs.nvidia.com/deeplearning/performance/dl-performance-fully-connected/index.xhtml#batch-size。另一个需要记住的事项是,你可能需要随着批次大小的增加而按比例增加学习率。

  • 使用融合优化器(如 FusedAdam 优化器)通过操作融合加速权重更新——将多个操作合并为一个。确保确认你的深度学习框架和硬件支持融合优化器。

这些是几种可能提高训练任务效率的常见参数。请注意,在许多实际应用场景中,你可能需要根据模型或任务的特定调优参数。

总结

在本章中,我们重点介绍了如何设计大规模数据并行、模型并行和混合分布式训练任务。我们讨论了如何根据具体的使用场景和模型架构选择并行类型。然后,我们回顾了几种常见的分布式训练组织方法——如参数服务器和 Allreduce 算法——以及调优分布式训练任务的各种性能考虑因素。现在,你将能够选择正确的分布式训练类型、技术栈和调试及调优训练任务性能的方法。接着,我们回顾了在 Amazon SageMaker 中使用流行的开源和专有库 SDDP 和 SMDP 进行分布式训练任务的几个示例。

运行大规模训练任务不仅需要初步的工程工作,还需要对训练任务进行良好的运营管理。在许多情况下,训练任务可能会持续数天甚至数周,或者你可能需要定期在新数据上重新训练模型。由于每个长期运行的深度学习训练任务都需要大量的计算资源和相应的时间及成本资源,因此我们希望确保训练过程高效。例如,我们需要实时控制模型在训练过程中是否正在收敛。否则,我们可能需要更早地停止训练,避免浪费计算资源和时间。在下一章中,我们将重点介绍如何为你的深度学习训练任务搭建一个运营栈。

第七章:深度学习训练的运营化

第一章《使用 Amazon SageMaker 介绍深度学习》中,我们讨论了 SageMaker 如何与 CloudWatch Logs 和 Metrics 集成,通过收集训练日志和指标来提供对训练过程的可视化。然而,深度学习DL)训练作业容易遇到与模型架构和训练配置相关的多种特定问题。需要专门的工具来监控、检测和应对这些问题。由于许多训练作业需要在大量计算实例上运行数小时甚至数天,因此错误的成本非常高。

在运行深度学习训练作业时,你需要意识到两种类型的问题:

  • 模型和训练配置的问题,阻碍了模型在训练过程中的高效学习。例如,梯度消失和爆炸、过拟合和欠拟合、损失未下降等问题。找出这些错误的过程被称为调试

  • 次优的模型和训练配置,未能充分利用可用的硬件资源。例如,假设批量大小小于最佳值,GPU 资源未得到充分利用,这导致训练速度比可能的速度慢。我们称这种找出问题的过程为分析

在本章中,我们将回顾用于训练、调试和分析的开源工具和 SageMaker 功能。我们将从流行的开源训练监控和调试工具TensorBoard开始,回顾它如何与 SageMaker 的训练基础设施集成。然后,我们将其与专有的SageMaker 调试器进行对比,后者提供了先进的功能,帮助你自动检测各种问题,并相应地管理训练作业。你将获得使用这两种工具的实际经验。

在运营化深度学习模型时,你通常需要解决的另一类问题是建立一种高效的方法来寻找最佳的模型超参数组合。这个过程被称为超参数调优。它在模型开发和采用的初期阶段尤为重要,因为此时你需要建立一个可以投入生产的模型基线。SageMaker 提供了一种自动化的方式,通过自动模型调优功能来调节你的模型。

最后,我们将讨论如何通过使用EC2 Spot 实例来降低训练作业和模型调优作业的成本。

在本章中,我们将涵盖以下主题:

  • 调试训练作业

  • 分析你的深度学习训练

  • 超参数优化

  • 使用 EC2 Spot 实例

阅读本章后,你将能够为大规模深度学习训练建立分析和调试程序,从而最小化不必要的成本和训练时间。你还将学会如何组织超参数调优,并利用 Spot 实例优化成本。

技术要求

在本章中,我们将提供代码示例,以便你能够培养实际技能。完整的代码示例可以在这里找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter7/

要跟随本代码进行操作,你将需要以下内容:

  • 一个具有管理 Amazon SageMaker 资源权限的 AWS 账户和 IAM 用户。

  • 配置一个 SageMaker 笔记本、SageMaker Studio 笔记本,或建立一个本地的 SageMaker 兼容环境。

  • 你需要访问 AWS 账户中的 GPU 训练实例。本章中的每个示例都会提供推荐的实例类型。你可能需要增加 SageMaker 训练作业 的计算配额,以启用 GPU 实例。在这种情况下,请按照docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.xhtml中的说明操作。

  • 你必须通过运行pip install -r requirements.txt来安装所需的 Python 库。包含所需库的文件位于 chapter7 目录的根目录下。

调试训练任务

为了有效地监控和调试深度学习训练任务,我们需要访问以下信息:

  • 标量值,如准确率和损失,用于衡量训练过程的质量

  • 张量值,如权重、偏差和梯度,代表模型及其优化器的内部状态

TensorBoard 和 SageMaker Debugger 都允许你收集张量和标量,因此两者都可以用于调试模型和训练过程。然而,不同于主要用于训练可视化的 TensorBoard,SageMaker Debugger 提供了几乎实时响应模型状态变化的功能。例如,如果训练损失在一段时间内没有下降,它可以让我们提前停止训练任务。

在本节中,我们将深入探讨如何使用 TensorBoard 和 SageMaker Debugger。我们将详细回顾这两种解决方案的功能,然后开发使用这两种解决方案调试训练脚本的实际经验。

请注意,我们将在调试和性能分析任务中使用相同的示例。

在 SageMaker 中使用 TensorBoard

TensorBoard 是一个最初为 TensorFlow 框架开发的开源工具,但现在也支持其他深度学习框架,包括 PyTorch。TensorBoard 支持以下功能,用于可视化和检查训练过程:

  • 随时间追踪标量值(损失、准确率等)。

  • 捕获张量,如权重、偏差和梯度,以及它们随时间变化的情况。这对于可视化权重和偏差并验证它们是否按预期变化非常有用。

  • 通过超参数仪表板进行实验追踪。

  • 将高维嵌入投影到低维空间。

  • 捕获图像、音频和文本数据。

此外,TensorBoard 还提供了针对 TensorFlow 程序的本地性能分析功能。通过附加组件,PyTorch 也支持性能分析。

调试 PyTorch 训练

让我们回顾一下 TensorBoard 如何帮助你深入了解训练过程,并通过实际示例调试它。我们将使用来自 PyTorch 模型库的预训练 ResNet 模型,并训练它识别两种类别:蜜蜂和蚂蚁。

我们在本节中提供了代码亮点。完整的训练代码可以在这里查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter7/1_TensorBoard_PyTorch.ipynb

修改训练脚本

要使用 TensorBoard,我们只需对训练脚本做最小的修改。请按照以下步骤操作:

  1. 首先,我们必须导入并初始化 TensorBoard 的 SummaryWriter 对象。在这里,我们使用 S3 位置来写入 TensorBoard 汇总:

    from torch.utils.tensorboard import SummaryWriter
    tb_writer = SummaryWriter(args.tb_s3_url)
    
  2. 接下来,我们必须捕获一些在训练过程中不会改变的训练文物——在我们这个案例中是模型图。请注意,我们需要在样本数据上执行模型的前向传播才能做到这一点:

    sample_inputs, _ = next(iter(dataloaders_dict["val"]))
    tb_writer.add_graph(model, sample_inputs, verbose=False, use_strict_trace=False)
    
  3. 在我们的训练循环中,我们捕获了我们希望检查的标量和张量。我们使用 epoch 编号作为时间维度。假设在我们这个案例中,我们希望捕获以下数据:

    • 每个 epoch 对于训练集和验证集的准确率和损失变化

    • 在训练阶段,第一个卷积层和最后一个全连接层的梯度和权重分布

    • 训练超参数以及它们对性能的影响

为了捕获这些参数,我们必须在训练循环中添加以下代码:

tb_writer.add_histogram("conv1.weight", model.conv1.weight, epoch)
tb_writer.add_histogram("conv1.weight_grad", model.conv1.weight.grad, epoch)
tb_writer.add_histogram("fc.weight", model.fc.weight, epoch)
tb_writer.add_histogram("fc.weight_grad", model.fc.weight.grad, epoch)
tb_writer.add_scalar(f"Loss/{phase}", epoch_loss, epoch)
tb_writer.add_scalar(f"Accuracy/{phase}", epoch_accuracy, epoch)
tb_writer.add_hparams(hparam_dict=vars(args), metric_dict={
                    f"hparam/loss_{phase}": epoch_loss,
                    f"hparam/accuracy_{phase}": epoch_accuracy})

现在,让我们回顾一下启用了调试的训练任务配置。

监控训练过程

要启动 SageMaker 训练任务,我们需要提供 TensorBoard 汇总文件将写入的 S3 位置。我们可以通过设置 tb-s3-url 超参数来实现,如下所示:

instance_type = 'ml.p2.xlarge'
instance_count = 1
job_name = "pytorch-tb-profiling-12"
tb_debug_path = f"s3://{bucket}/tensorboard/{job_name}"
estimator = PyTorch(
          entry_point="train_resnet_tb.py",
          source_dir='1_sources',
          role=role,
          instance_type=instance_type,
          sagemaker_session=sagemaker_session,
          image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.10.2-gpu-py38-cu113-ubuntu20.04-sagemaker",
          instance_count=instance_count,
          hyperparameters={
              "batch-size":64,
              "num-epochs":10,
              "input-size" : 224,
              "feature-extract":False,
              "tb-s3-url": tb_debug_path,
              "num-data-workers": 4
          },
          disable_profiler=True,
          debugger_hook_config=False,
          base_job_name=job_name,
      )

训练任务开始后,你可以通过在终端中运行以下命令,启动本地的 TensorBoard:

tensorboard --logdir ${tb_debug_path} 

在云开发环境中使用 TensorBoard 时,请注意以下事项:

  • 如果你使用的是 SageMaker Notebook 实例,则可以通过以下地址访问 TensorBoard:https://YOUR_NOTEBOOK_DOMAIN/proxy/6006/

  • 如果你使用的是 SageMaker Studio,则可以通过以下地址访问 TensorBoard:https://<YOUR_STUDIO_DOMAIN>/jupyter/default/proxy/6006/

随着训练任务的进展,TensorBoard 数据将实时更新。让我们在 TensorBoard 中回顾我们的训练过程:

  • 标量(Scalar)时间序列(Time Series)标签页中,您可以看到标量值随时间变化的情况。我们使用 epoch 索引作为时间的指示器。图 7.1 显示了每个 epoch 的训练和验证准确率:

图 7.1 – TensorBoard 中随时间变化的准确率

图 7.1 – TensorBoard 中随时间变化的准确率

  • 图形(Graph)标签页中,您可以看到模型的可视化表示,以及数据从输入到输出的流动方式。

  • 0,表示我们的模型正在学习,因此绝对梯度值在下降:

图 7.2 – TensorBoard 中模型权重的直方图

图 7.2 – TensorBoard 中模型权重的直方图

  • HParam 标签页使我们能够并排捕捉和比较超参数。这对于在超参数搜索过程中跟踪实验,识别最优的模型和训练任务配置非常有用。

现在我们已经理解了如何使用 TensorBoard 可视化训练过程,让我们来看看如何使用 TensorBoard 对训练任务进行性能分析。

性能分析 PyTorch 训练

TensorBoard 为 TensorFlow 程序(包括 Keras)提供了开箱即用的性能分析功能。要在 TensorBoard 中对 PyTorch 程序进行性能分析,您可以使用开源的torch_tb_profiler插件。

在进行训练过程的性能分析时,我们通常关心以下几个方面:

  • 我们如何高效地利用资源(GPU 和 CPU)随时间变化的情况

  • 哪些操作(深度学习运算符、数据加载、内存传输等)使用了哪些资源

  • 在分布式训练的情况下,节点与各个训练设备之间的通信效率如何

  • 如何提高整体资源利用率并提高训练效率

TensorFlow 和 PyTorch 插件都为 TensorBoard 提供了性能分析的功能。让我们回顾一下性能分析是如何与调试任务一起工作的。

修改训练脚本

要使用 torch_tb_profiler 对应用程序进行性能分析,我们需要对训练代码进行最小的修改。具体来说,我们需要用插件上下文管理器包裹训练循环,如下代码块所示:

with torch.profiler.profile(
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=5),
    on_trace_ready=torch.profiler.tensorboard_trace_handler(
        os.path.join(os.environ["SM_OUTPUT_DATA_DIR"], "tb_profiler")
    ),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for _, (inputs, labels) in enumerate(dataloaders[phase]):
      # The rest of training loop without changes

初始化时传递给上下文管理器的参数定义了必须收集哪些性能分析数据以及在什么时间间隔收集。在写这本书时,torch_db_profiler 插件不支持写入 S3 位置。因此,我们必须将性能分析数据写入存储在 "SM_OUTPUT_DATA_DIR" 环境变量中的本地输出目录。训练完成后,SageMaker 会自动将该目录的内容归档并存储到 S3 位置。

使用 TensorBoard Profiler

为了查看 TensorBoard Profiler 的输出,我们需要将数据下载到本地环境中:

  1. 我们将从获取性能分析数据的路径开始。为此,我们可以使用训练任务估算器实例:

    tb_profiler_path = f"{estimator.latest_training_job.describe()['OutputDataConfig']['S3OutputPath']}{estimator.latest_training_job.describe()['TrainingJobName']}/output/output.tar.gz"
    
  2. 然后,在你的笔记本或终端窗口中,你可以运行以下命令来解压分析器数据并启动 TensorBoard:

    aws s3 cp ${ tb_profiler_path} .
    mkdir profiler_output
    tar -xf output.tar.gz -C profiler_output
    tensorboard --logdir ./profiler_output
    

启动 TensorBoard 后,你应该会自动跳转到分析器摘要页面。在这里,你可以访问几个包含分析信息的视图:

  • Overview 标签页提供了用于训练的设备(设备)的概览,展示了它们的时间利用率和操作的拆解。例如,在我们的案例中,大部分时间都花费在执行包括前向和反向模型传递的内核上。这通常是我们在训练模型时充分利用 GPU 资源的一个良好指标:

图 7.3 – TensorBoard Profiler 的 Overview 标签页

图 7.3 – TensorBoard Profiler 的 Overview 标签页

  • Operators 标签页让你了解特定操作符(如卷积或批量归一化)消耗了多少时间。在下图中,我们可以看到,例如,卷积层的反向传递占用了大部分 GPU 时间:

图 7.4 – TensorBoard Profiler 的 Operators 标签页

图 7.4 – TensorBoard Profiler 的 Operators 标签页

  • Kernel 标签页显示了在特定 GPU 核心上执行的时间。比如在下图中,你可以看到各种 单精度通用矩阵乘法SGEMM)内核占用了大部分时间:

图 7.5 – TensorBoard Profiler 的 Kernel 标签页

图 7.5 – TensorBoard Profiler 的 Kernel 标签页

  • Trace 标签页显示了分析的操作符和 GPU 内核的时间线,以及 CPU 和 GPU 设备之间的交接(例如,将数据输入从 CPU 转移到 GPU):

图 7.6 – TensorBoard Profiler 的 Trace 标签页

  • Memory 标签页提供了给定设备的内存利用情况。在下图中,你可以看到分配的内存(即用于存储张量的内存)和总保留内存:

图 7.7 – TensorBoard Profiler 的 Memory 标签页

图 7.7 – TensorBoard Profiler 的 Memory 标签页

如你所见,TensorBoard 是一个非常适合用于监控、调试和分析训练脚本的工具。当与插件一起使用时,TensorBoard 支持 TensorFlow 和 PyTorch 框架。然而,TensorBoard 的一个缺点是,它没有提供任何方式来响应不理想的情况,比如 GPU 设备未充分利用,或模型在训练过程中出现收敛缓慢或无收敛的情况。为了在这种情况下提前停止训练作业,你需要通过回调和自定义逻辑进一步对代码进行工具化。

SageMaker Debugger 通过提供一种通用机制来检测常见的训练问题并采取缓解措施,解决了这些局限性。

使用 SageMaker Debugger 监控训练

SageMaker Debugger 是一个全面的 SageMaker 功能,允许你自动监控、调试和分析运行在 SageMaker 上的深度学习训练任务。SageMaker Debugger 通过捕捉训练循环的内部状态和实例指标,提供近实时的深度学习训练洞察。Debugger 还可以帮助你自动检测训练过程中常见问题,并在发现问题时采取适当的行动。这使得你能够在复杂的深度学习训练任务中更早地发现问题,并作出相应反应。此外,SageMaker Debugger 还支持编写自定义规则,以应对内置规则未涵盖的场景。

SageMaker 具有几个关键组件:

  • 开源的smedebug库(https://github.com/awslabs/sagemaker-debugger),它与深度学习框架和 Linux 实例集成,将调试和分析数据持久化到 Amazon S3,并在训练任务启动后检索和分析数据

  • SageMaker Python SDK,允许你在训练脚本中通过最小的代码变更来配置smedebug

  • 自动化配置处理任务,以验证输出张量和分析数据是否符合规则

SageMaker Debugger 支持 TensorFlow、PyTorch 和 MXNet 深度学习框架。smedebug库默认安装在 SageMaker 深度学习容器中,因此你可以在不修改训练脚本的情况下开始使用 SageMaker Debugger。你还可以在自定义 Docker 容器中安装smdebug库,并使用 SageMaker Debugger 的所有功能。

注意

请注意,不同深度学习框架的smedebugAPI 可能会有细微的差异。

smedebug库提供了丰富的 API,用于配置、保存和分析捕获的张量。它通过在训练过程中注入hook对象来捕获张量和标量。hook允许你将张量和标量分组到逻辑张量Trial对象中,从而可以查询给定训练任务的存储张量进行进一步分析。你可以实时运行张量查询,而无需等待训练任务完全完成。SageMaker Debugger 还支持生成兼容 TensorBoard 的摘要日志,方便可视化输出张量和标量。

使用 SageMaker Debugger

让我们将这些概念应用于一个实际任务。我们将对 ResNet 模型进行工具化,并针对二分类任务进行微调。完整代码请查看此处:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter7/2_SMDebugger_PyTorch.ipynb

代码工具化

smedebug库只需要最少的修改即可捕获张量和标量。首先,你需要在训练循环外部初始化hook对象,以及在模型和优化器初始化之后:

...
model = initialize_resnet_model(
    NUM_CLASSES, feature_extract=False, use_pretrained=True
)
model.to(torch.device("cuda"))
optimizer = optim.SGD(params_to_update, lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
hook = smd.Hook.create_from_json_file()
hook.register_hook(model)
hook.register_loss(criterion)
...

请注意,我们正在使用 .create_from_json_file() 方法来创建我们的 hook 对象。此方法基于您在 SageMaker 训练对象中提供的 hook 配置实例化 hook。由于我们将 modelcriterion 对象都添加到 hook 中,因此我们应该期望看到模型参数(权重、偏置等),以及损失标量。

在我们的训练循环中,我们唯一需要修改的地方是通过在 smedebug.modes.Trainsmedebug.modes.Eval 之间切换来区分训练阶段和验证阶段。这将使得 smedebug 能够区分在训练和评估阶段捕获的张量:

  for epoch in range(1, args.num_epochs + 1):
        for phase in ["train", "val"]:
            if phase == "train":
                model.train()  # Set model to training mode
                if hook:
                    hook.set_mode(modes.TRAIN)
            else:
                model.eval()  # Set model to evaluate mode
                if hook:
                    hook.set_mode(modes.EVAL)
            running_corrects = 0
            running_loss = 0.0
            step_counter = 0
            epoch_start = time.time()
            for _, (inputs, labels) in enumerate(
            dataloaders[phase]):
            # inside training loop
...

现在,让我们回顾一下在运行 SageMaker 训练任务时,如何配置 hook、规则、操作和张量集合。

训练任务配置

作为 SageMaker Python SDK 的一部分,AWS 提供了 sagemaker.debugger 库用于 Debugger 配置。让我们来看看:

  1. 我们将首先导入一些 Debugger 实体:

    from sagemaker.debugger import (
        Rule,
        DebuggerHookConfig,
        TensorBoardOutputConfig,
        CollectionConfig,
        rule_configs,
        ProfilerRule
    )
    
  2. 然后,我们必须定义自动化操作和一组规则。在这里,我们使用 Debugger 内置的规则来检测一些常见的深度学习训练问题。请注意,我们可以为不同的规则分配不同的操作。在我们的例子中,我们希望在触发规则时立即停止训练任务:

    actions = rule_configs.ActionList(
        rule_configs.StopTraining())
    rules = [
        Rule.sagemaker(rule_configs.vanishing_gradient(), actions=actions),
        Rule.sagemaker(rule_configs.overfit(), actions=actions),
        Rule.sagemaker(rule_configs.overtraining(), actions=actions),
        Rule.sagemaker(rule_configs.poor_weight_initialization(), actions=actions),
    ]
    
  3. 接下来,我们必须配置张量的收集以及它们的持久化方式。在这里,我们将定义要持久化权重和损失的集合。对于权重,我们还将保存一个可以在 TensorBoard 中进一步可视化的直方图。我们还将为训练和评估阶段设置保存间隔:

    collection_configs=[
            CollectionConfig(
                name="weights",
                parameters={
                    "save_histogram": "True"
                    }
                ),
            CollectionConfig(name="losses"),
        ]
    hook_config = DebuggerHookConfig(
        hook_parameters={"train.save_interval": "1", "eval.save_interval": "1"},
        collection_configs=collection_configs
    )
    
  4. 现在,我们准备将这些对象传递给 SageMaker 的 Estimator 对象:

    tb_debug_path = f"s3://{bucket}/tensorboard/{job_name}"
    tensorboard_output_config = TensorBoardOutputConfig(
        s3_output_path=tb_debug_path
    )
    debug_estimator = PyTorch(
              entry_point="train_resnet_sm.py",
              source_dir='2_sources',
              role=role,
              instance_type=instance_type,
              sagemaker_session=sagemaker_session,
              image_uri=image_uri,
              instance_count=instance_count,
              disable_profiler=True,
              rules=rules,
              debugger_hook_config=hook_config,
              tensorboard_output_config=tensorboard_output_config,
              base_job_name=job_name,
          )
    

现在,我们准备使用 fit() 方法开始训练任务。在下一节中,我们将学习如何检索和分析 SageMaker Debugger 输出。

审查 Debugger 结果

SageMaker Debugger 提供了一个功能,可以从训练任务中检索和分析收集的张量,作为 smedebug 库的一部分。在接下来的步骤中,我们将介绍一些关键的 API:

  1. 在下面的代码块中,我们使用保存张量的 S3 路径来创建一个新的 trial 对象:

    import smdebug.pytorch as smd
    tensors_path = debug_estimator.latest_job_debugger_artifacts_path()
    trial = smd.create_trial(tensors_path)
    
  2. 现在,让我们通过运行以下命令输出所有可用的张量:

    print(f"Persisted tensors: {trial.tensor_names()}")
    
  3. 您应该能够看到多个集合,其中包含许多张量,包括偏置、权重、损失和梯度。让我们访问具体的数值。运行以下命令将返回一组相关的标量值:

    print(f"Loss values {trial.tensor('CrossEntropyLoss_output_0').values()}")
    
  4. 使用一个简单的绘图函数(有关其实现,请参考源代码),我们可以可视化训练和评估阶段的损失。运行以下命令将生成一个 2D 损失图表。类似地,您可以访问和处理张量:

    plot_tensor(trial, "CrossEntropyLoss_output_0")
    

下图可视化了训练和验证损失:

图 7.8 – 训练和验证损失

图 7.8 – 训练和验证损失

  1. 现在,让我们回顾一下在训练过程中是否触发了任何规则:

    for s in debug_estimator.latest_training_job.rule_job_summary():
        print(f"Rule: {s['RuleConfigurationName']}",
              f"status: {s['RuleEvaluationStatus']}")
    

这将输出所有已配置的规则;它们的状态如下:

Rule: VanishingGradient, status: NoIssuesFound
Rule: Overfit, status: NoIssuesFound 
Rule: Overtraining, status: NoIssuesFound 
Rule: PoorWeightInitialization, status: NoIssuesFound

如我们所见,在我们的案例中,没有触发任何规则,工作已经完成。你可以尝试不同的规则设置。例如,你可以重置模型层中的某一层权重,这将触发PoorWeightInitiailization规则,进而导致训练过程被停止。

  1. 最后,让我们使用 TensorBoard 对保存的张量进行可视化检查。为此,我们只需使用之前提供给Estimator对象的 S3 路径启动 TensorBoard:

    ! tensorboard --logdir  {tb_debug_path}
    

你可以随意探索 TensorBoard。你应该会看到权重的直方图。

在本节中,我们回顾了 SageMaker Debugger 的关键功能,并学习了如何使用它们。你可能已经注意到 SageMaker Debugger 相比于 TensorBoard 的一些优势:

  • 在为 SageMaker Debugger 进行代码仪器化时几乎无需任何额外工作

  • 提供强大的 API 来处理和分析输出张量

  • 大量内置规则和操作,支持创建自定义规则和操作

  • TensorBoard 功能开箱即用

通过这些功能,SageMaker Debugger 允许你提高训练工作的质量,加速实验,并减少不必要的成本。

此外,SageMaker Debugger 还提供了性能分析功能。我们接下来将回顾它们。

对你的深度学习训练进行性能分析

SageMaker Debugger 允许你从训练实例中收集各种类型的高级指标。一旦这些指标被收集,SageMaker 将生成详细的指标可视化,检测资源瓶颈,并提供如何提高实例利用率的建议。

SageMaker Debugger 收集两种类型的指标:

  • 系统指标:这些是训练实例的资源利用率指标,如 CPU、GPU、网络和 I/O。

  • 框架指标:这些指标是在深度学习框架级别收集的,包括原生框架分析器(如 PyTorch Profiler 或 TensorFlow Profiler)收集的指标、数据加载器指标和 Python 性能分析指标。

与调试相似,你可以定义一些规则,自动评估收集的指标。如果触发了某个规则,你可以定义一个或多个将采取的操作。例如,如果训练任务的 GPU 利用率低于某个阈值,你可以发送一封电子邮件。

现在是时候使用 SageMaker Debugger 对我们的训练代码进行性能分析了。你可以在 https://github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter7/2_SMDebugger_PyTorch.ipynb 中的Profiling DL Training部分找到完整的代码。

配置训练任务以进行性能分析

我们将首先定义要收集的系统和框架指标。例如,我们可以为框架、数据加载器和 Python 提供自定义配置。请注意,系统性能分析默认是启用的:

from sagemaker.debugger import (ProfilerConfig, 
                                FrameworkProfile, 
                                DetailedProfilingConfig, 
                                DataloaderProfilingConfig, 
                                PythonProfilingConfig,
                                PythonProfiler, cProfileTimer)
profiler_config=ProfilerConfig(
    system_monitor_interval_millis=500,
    framework_profile_params=FrameworkProfile(
        detailed_profiling_config=DetailedProfilingConfig(
            start_step=2, 
            num_steps=1),
        dataloader_profiling_config=DataloaderProfilingConfig(
            start_step=2, 
            num_steps=1),
        python_profiling_config=PythonProfilingConfig(
            start_step=2, 
            num_steps=1, 
            python_profiler=PythonProfiler.CPROFILE, 
            cprofile_timer=cProfileTimer.TOTAL_TIME)))

然后,我们必须将分析配置提供给 SageMaker 训练作业配置:

profiler_estimator = PyTorch(
          entry_point="train_resnet_sm.py",
          source_dir='2_sources',
          role=role,
          instance_type='ml.p2.xlarge',
          sagemaker_session=sagemaker_session,
          image_uri=image_uri,
          instance_count=instance_count,
          hyperparameters={
              "num-data-workers":8,
          },
          disable_profiler=False,
          profiler_config=profiler_config,
          rules=rules,
        #  debugger_hook_config=hook_config,
        #  tensorboard_output_config=tensorboard_output_config,
          base_job_name=job_name,
      )

请注意,我们将 num-data-workers 设置为 8,而 ml.p2.xlarge 只有 4 个 CPU 核心。通常建议数据工作者的数量与 CPU 数量相等。让我们看看 SageMaker Debugger 是否能够检测到这个次优配置。

审查分析结果

您可以开始实时监控分析结果。我们将使用 semdebug.profiler API 处理分析输出:

training_job_name = profiler_estimator.latest_training_job.job_name
region = "us-east-1"
tj = TrainingJob(training_job_name, region)
tj.wait_for_sys_profiling_data_to_be_available()

一旦数据可用,我们可以提取并可视化它。运行以下代码将绘制来自系统指标的 CPU、GPU 和 GPU 内存利用率图:

from smdebug.profiler.analysis.notebook_utils.timeline_charts import TimelineCharts
system_metrics_reader = tj.get_systems_metrics_reader()
system_metrics_reader.refresh_event_file_list()
view_timeline_charts = TimelineCharts(
    system_metrics_reader,
    framework_metrics_reader=None,
    select_dimensions=["CPU", "GPU"],
    select_events=["total"],
)

同样,您也可以可视化其他收集的指标。SageMaker Debugger 还会生成一份详细的分析报告,将所有可视化内容、见解和建议集中在一个地方。一旦您的训练作业完成,您可以通过在终端运行以下命令下载分析报告和所有收集的数据:

aws s3 cp s3://<JOB_BUCKET>/<JOB_NAME>/rule-output ./ --recursive

一旦所有资产都已下载,请在浏览器中打开 profiler-report.xhtml 文件,并查看生成的信息。或者,您也可以打开 profiler-report.ipynb,它以可执行的 Jupyter notebook 形式提供相同的见解。

报告涵盖以下方面:

  • 系统使用统计信息

  • 框架指标总结

  • 规则及其状态总结

  • 训练循环分析和优化建议

请注意,在 数据加载分析 部分,您应该看到根据我们的预期减少数据工作者数量的建议。

如您所见,SageMaker Debugger 提供了广泛的分析功能,包括建议改善和自动化规则验证,且开发工作量最小。与其他 Debugger 功能类似,只要使用内置规则,分析是免费的。

超参数优化

SageMaker 自动模型调优作业允许您并行运行多个训练作业,每个作业使用唯一的超参数组合。换句话说,一个调优作业会创建多个 SageMaker 训练作业。超参数调优通过并行尝试多个超参数组合并迭代地朝着更优的组合前进,可以加速模型开发和优化。然而,它并不保证您的模型性能总是会提高。例如,如果所选的模型架构不适合当前任务,或者您的数据集对于所选模型来说太小,那么在进行超参数优化时,您不太可能看到任何改进。

在设计调优作业时,您需要考虑几个关键参数,如下所示:

  • 搜索算法(或 策略):这定义了 SageMaker 如何选择下一个超参数组合。

  • 具有范围的超参数:SageMaker 搜索算法将在用户定义的范围内选择超参数值。

  • 目标指标:这将用于比较一组超参数并定义最佳候选值。SageMaker 不限制你选择任何任意的目标指标。

SageMaker 支持两种搜索策略:贝叶斯随机。随机搜索在定义的范围内随机选择下一组超参数。尽管这是一种简单的策略,但被认为是相对高效的。因为下一组超参数的选择不依赖于之前尝试过的或当前运行的组合,所以你可以同时运行大量的训练作业。贝叶斯搜索则是基于之前训练作业的结果来选择下一组超参数。在后台,SageMaker 为此训练了一个回归模型,该模型将之前作业的结果作为输入(超参数及其对应的目标指标),并输出候选的超参数组合。需要注意的是,贝叶斯模型可能不会收敛。在这种情况下,回顾已确定的超参数范围是有意义的。

选择超参数及其范围对调优作业的性能有显著影响。SageMaker 支持几种类型的超参数——分类的、连续的和整数型的。你可以组合不同类型的超参数。例如,以下代码将模型架构定义为分类超参数,学习率调度步长定义为整数型参数,而学习率定义为连续型参数(换句话说,浮动类型):

hyperparameter_ranges = {
"model_type" : sagemaker.tuner.CategoricalParameter(["resnet", "vgg16", "densenet"]}]),
"learning_rate" : sagemaker.tuner.ContinuousParameter(0.0001,0.1, scaling_type="Logarithmic"),
"lr_scheduler_step_size" : sagemaker.tuner.IntegerParameter(10,100, scaling_type="Linear"),
}

请注意,对于数值型超参数,我们还为学习率参数定义了“对数”缩放类型,因为它的范围跨越了多个数量级。对于调度步长,我们选择了“线性”缩放类型,因为它的范围较窄。

你还需要为超参数调优作业定义目标指标。目标指标的定义类似于其他指标,通过正则表达式模式定义。请注意,你需要确保训练脚本将目标指标输出到 stdout/stderr 流中。请按以下步骤操作:

  1. 在以下代码中,我们定义了四个指标,这些指标将由 SageMaker 捕获,然后选择 val_accuracy 作为我们要优化的目标指标:

    metric_definitions = [
        {"Name": "train_loss",
         "Regex": "Train Loss = (.*?);"},
        {"Name": "val_loss",
         "Regex": "Val Loss=(.*?);"},
        {"Name": "train_accuracy",
         "Regex": "Train Accuracy = (.*?);"},
        {"Name": "val_accuracy",
         "Regex": "Val Accuracy = (.*?);"},]
    objective_metric_name = "val_accuracy"
    
  2. 接下来,我们需要定义训练作业的参数。请注意,作为训练作业配置一部分提供的超参数将是静态的,并且在调优作业中不会发生变化:

    estimator = PyTorch(
              entry_point="train_script.py",
              role=role,
              instance_type=instance_type,
              sagemaker_session=sagemaker_session,
              image_uri=image_uri,
              instance_count=instance_count,
              hyperparameters={
                  "batch-size":64,
                  "num-epochs":5,
              })
    
  3. 然后,我们必须在 HyperParameterTuner 对象中结合我们的目标指标、指标定义和超参数范围,该对象将协调创建子训练作业并跟踪调优的整体状态。此外,我们还必须提供训练作业的最大总数和并发训练作业的数量。这些参数将影响调优作业的运行速度和总成本:

    tuner = sagemaker.tuner.HyperparameterTuner(
         estimator,
         objective_metric_name,
         hyperparameter_ranges,
         metric_definitions,
         objective_type="Maximize",
         max_jobs=200,
         max_parallel_jobs=10)
    

此外,注意 objective_type 参数,它定义了调优作业是要最大化还是最小化目标指标。由于我们选择了 accuracy 作为目标指标,因此我们希望最大化它。

  1. 一旦调优对象被实例化,您可以使用 .fit() 方法开始训练:

    tuner.fit({"train":train_data_location, "val":val_data_location})
    
  2. 作业完成后,您可以分析调优作业的结果。为此,您可以导航到 AWS 控制台并进行可视化检查。或者,您可以将调优作业结果和统计数据导出为 pandas DataFrame 以进行进一步分析,如下所示:

    tuner = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)
    tuner_results = tuner.dataframe()
    

使用此方法,您可以执行更高级的分析,例如定义各种超参数与目标指标之间的相关性。例如,这种类型的分析可能揭示出需要修改超参数范围的情况,从而进一步提升目标指标。

使用 EC2 Spot 实例

运行大规模的训练和模型调优作业可能非常昂贵。最小化成本的一种方法是使用来自选定 AWS 区域的未使用计算资源池中的 EC2 Spot 实例。因此,Spot 实例比常规按需实例便宜得多(最高可达 90%)。然而,如果所选实例类型在给定 AWS 区域的 Spot 容量耗尽,Spot 实例可能会被随时停止。

SageMaker 简化了为训练作业提供 Spot 实例的过程,并在 Spot 容量再次可用时完全处理中断和训练作业重启。当训练作业被中断并重新启动时,我们希望继续训练过程,而不是从头开始。为了支持这一点,您的训练脚本需要进行修改,以便能够保存并重新启动训练作业。

为了支持 Spot 训练,您的训练脚本需要以下修改:

  • 当首次加载模型时,检查 /opt/ml/checkpoints 路径中是否已有模型副本。如果检查点模型已存在,说明我们之前训练过此模型。要继续训练,我们需要加载检查点模型并继续训练。如果检查点模型不存在,则继续常规的模型加载:

图 7.9 – 将检查点工件上传到 S3 存储

图 7.9 – 将检查点工件上传到 S3 存储

  • 在你的训练脚本中,你需要指定检查点处理程序(请参考深度学习框架的文档),并将模型检查点存储在指定的目录中——即 /opt/ml/checkpoints。如果 Spot 实例被中断,SageMaker 将自动将该目录的内容复制到 S3。Spot 实例恢复可用后,SageMaker 会从 S3 将检查点复制回 /opt/ml/checkpoints 目录:

图 7.10 – 从 S3 存储恢复检查点文件

图 7.10 – 从 S3 存储恢复检查点文件

在使用 Spot 实例时,请注意,使用 Spot 训练可能会导致训练时间变长且不可预测。每次 Spot 实例中断都会导致重启时额外的启动时间。可用的 Spot 容量取决于实例类型和 AWS 区域。在某些 AWS 区域,基于 GPU 的实例类型可能具有非常有限的 Spot 容量。请注意,Spot 容量会不断波动。你可以使用Amazon Spot 实例顾问功能,来确定不同 EC2 实例的可用 Spot 容量、中断的可能性以及与常规按需实例相比的成本节省。

总结

本章总结了本书的第二部分。在本章及前两章中,我们讨论了如何构建和优化大规模训练作业。首先,我们回顾了可用于深度学习训练的专业硬件,并介绍了如何选择最优的实例类型。接着,我们讨论了如何使用开源解决方案和 Amazon 专有解决方案来进行分布式训练。在本章中,我们讨论了如何高效地将模型训练操作化。我们回顾了训练过程中可能遇到的不同问题,以及如何检测和缓解这些问题。我们还讨论了如何管理和优化超参数调优。

第三部分提供深度学习模型服务中,我们将深入探讨在 Amazon SageMaker 上进行深度学习推理的过程。我们将讨论可用于推理的硬件,以及如何设计你的推理服务器。然后,我们将回顾模型服务的操作方面。在下一章,第八章考虑推理硬件中,我们将回顾适用于推理工作负载的硬件加速器,讨论选择标准,并解释如何使用模型编译器和 SageMaker Neo 对你的模型进行优化,以便在特定硬件加速器上进行推理。

第三部分:提供深度学习模型服务

在本章中,我们将重点讨论如何在 Amazon SageMaker 上托管训练好的模型。我们将回顾可用的软件和硬件选项,并提供选择建议以及何时使用它们的指导。

本节包含以下章节:

  • 第八章考虑推理硬件

  • 第九章实现模型服务器

  • 第十章操作化推理工作负载

第八章:考虑用于推理的硬件

在本书的第三部分,深度学习模型的服务中,我们将重点讨论如何开发、优化和将推理工作负载实现生产化,适用于深度学习DL)模型。与训练类似,DL 推理是计算密集型的,且需要理解特定类型的硬件,这些硬件专为推理设计,此外还需要掌握模型优化技术,以及专门的软件服务器来管理模型部署并处理推理流量。Amazon SageMaker 提供了广泛的功能来解决这些方面的问题。

在本章中,我们将讨论模型服务的硬件选项和模型优化。我们将回顾适用于 DL 推理的可用硬件加速器,并讨论如何选择其中之一。Amazon SageMaker 提供了多种 NVIDIA GPU 加速器和专为 DL 推理设计的专有芯片——AWS Inferentia。SageMaker 还允许您使用其 Elastic Inference 功能访问加速器容量。由于每个推理用例都是独特的,并且有其特定的业务需求,我们将提出一套选择标准,供您在评估最适合推理的硬件加速器时参考。

在构建 DL 推理工作负载时,另一个重要的方面是理解如何针对目标硬件加速器优化特定的模型架构。这一过程被称为模型编译。我们将回顾流行的优化器和运行时环境——NVIDIA TensorRT,它为在 NVIDIA GPU 加速器上运行的模型提供最佳的延迟和吞吐量。接着,我们将讨论 Neuron SDK,它优化模型以便在 AWS Inferentia 芯片上运行。我们还将讨论 SageMaker Neo——一种托管的模型编译服务,允许您为各种数据中心和边缘硬件加速器编译模型。请注意,本书不涉及任何边缘或嵌入式平台。

在本章中,我们将讨论以下主题:

  • 在 AWS 云中选择硬件加速器

  • 为推理编译模型

阅读完本章后,您将能够为您的推理工作负载选择一种高效的硬件配置,具有最佳的价格/性能特性,并进行进一步的优化。

技术要求

在本章中,我们将提供代码示例,帮助您培养实际技能。完整的代码示例可以在此处查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter8/

要跟随此代码,您需要以下内容:

  • 一个 AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户。

  • 已建立 SageMaker 笔记本、SageMaker Studio 笔记本或本地兼容 SageMaker 的环境。

  • 访问 AWS 账户中的 GPU 训练实例。本章中的每个示例将提供推荐的实例类型。你可能需要增加SageMaker 训练作业的计算配额,以启用 GPU 实例。在这种情况下,请按照docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.xhtml中的说明进行操作。

  • 你必须通过运行pip install -r requirements.txt来安装所需的 Python 库。包含所需库的文件可以在chapter8目录中找到。

  • 在本章中,我们将提供编译模型以进行推理的示例,这需要访问特定的加速器类型。如果你打算跟随这些代码示例,请准备好具有目标加速器的 SageMaker 笔记本实例或 SageMaker Studio 笔记本。

在 AWS Cloud 中选择硬件加速器

AWS Cloud 和 Amazon SageMaker 提供一系列适用于推理工作负载的硬件加速器。选择硬件平台通常需要通过多次实验,使用不同的加速器和服务参数进行测试。我们来看一些在评估过程中可能有用的关键选择标准。

延迟-吞吐量权衡

推理延迟定义了你的模型能多快地返回推理结果给最终用户,我们希望最小化延迟以改善用户体验。推理吞吐量定义了可以同时处理多少个推理请求,我们希望最大化吞吐量,以确保尽可能多的推理请求能够被处理。在软件工程中,通常会讨论延迟和吞吐量之间的权衡,因为通常很难同时最小化延迟并最大化吞吐量,因此你需要在这两个特性之间找到平衡。

在具体的使用案例中,通常会将目标延迟和吞吐量 SLA 作为业务需求的一部分。找到可接受的延迟和吞吐量的权衡需要通过与不同加速器以及模型/服务器参数进行基准测试,以符合目标 SLA。

例如,在运行实时推理端点时,通常关注的是延迟 SLA,因为它会直接影响到最终用户。你可以先找到延迟在目标 SLA 范围内的硬件和模型配置,然后再扩展以达到期望的吞吐量。在批量推理的情况下,整体系统吞吐量通常比延迟更为重要。我们希望最大化吞吐量,以确保我们的硬件资源得到有效利用。

成本

运行推理工作负载的成本是另一个重要参数,它会影响你使用的硬件以及延迟和吞吐量服务水平协议(SLA)。虽然 AWS 和 SageMaker 提供了市场上最强大的 GPU 加速器之一,但它们的成本对于你的特定用例可能过于高昂。因此,通常情况下,你可能需要调整延迟和吞吐量 SLA,以使你的深度学习推理应用在经济上可行。

支持的框架和操作符

运行推理工作负载需要大量的计算资源,因为它需要对单个或一批输入执行模型的前向传递。每次前向传递由一系列独立的计算任务组成。计算任务的一种类型被称为操作符。常见的深度学习操作符的例子包括矩阵乘法、卷积和平均池化。

深度学习框架支持的操作符始终可以在 CPU 设备上运行。然而,正如我们在 第五章 中讨论的,考虑用于训练的硬件,CPU 并不是最有效的深度学习加速器。因此,使用 CPU 运行推理会导致比使用 GPU 和 ASIC 芯片等专用加速器更高的延迟。

NVIDIA GPU 加速器通过 CUDA 工具包支持广泛的操作符。在某些情况下,你可能需要为你的特定模型架构实现新的操作符。CUDA 工具包为此类自定义操作符开发提供了编程 API。

AWS Inferentia 等 ASIC 加速器支持有限的操作符和框架列表。在某些操作符不受支持的情况下,该操作符将由 CPU 设备执行。这使得你可以在专用加速器上运行许多模型架构,但另一方面,它很可能会导致推理延迟增加,因为 CPU 执行的整体缓慢性以及在模型前向传递过程中 ASIC 和 CPU 加速器之间的必要数据传递。

因此,在选择目标硬件加速器时,你需要了解支持哪些深度学习框架和操作符。

第五章 中,考虑用于深度学习训练的硬件,我们概述了 Amazon SageMaker 平台上可用的深度学习硬件加速器。在接下来的部分,我们将重点介绍一些推荐用于推理工作负载的加速器和计算实例。

G4 实例系列 – 具有最佳价格和性能比的推理选项

G4 实例配备了 NVIDIA T4 Tensor Core GPU,拥有 16 GB 的内存。这款加速器是 NVIDIA 为云计算和数据中心的推理任务设计的。它支持 FP32、FP16、INT8 和 INT4 精度类型。由于 G4 综合了与推理工作负载相关的性能特征,并且与更强大的 P3 系列相比成本更低,因此应该被视为运行深度学习推理工作负载的默认选项。

为了进一步优化性能,你可以使用 NVIDIA TensorRT 优化器编译你的模型。我们将在下一节详细讨论 TensorRT 优化器。

P3 实例系列——适用于推理的高性能但昂贵

配备 NVIDIA V100 加速器的 P3 实例主要是为大规模训练设计的。与 G4 相比,P3 系列最多可提供 32 GB 的 GPU 内存和更大的网络带宽(包括 GPU 间和节点间的带宽)。P3 还支持 F64、FP32、FP16 和 INT8 精度类型。

P3 的许多特性非常适合大规模分布式训练,但对于推理则不太相关。例如,你很少需要使用双精度类型;相反,你希望在推理过程中降低精度,以减少延迟。更高的网络带宽(特别是节点间带宽)对于推理工作负载也不太相关,因为在推理时通常不需要将模型分布到不同节点上。

因此,虽然 P3 系列的性能优于 G4,但其成本更高,对于推理工作负载的好处很少。你可能想要选择 P3 而非 G4 的一种情况是,当你正在运行大模型的推理时。在这种情况下,P3dn.24xlarge 实例可以为你提供 8 个每个有 32 GB 内存的 V100 GPU。

重要提示

请注意,这里我们只考虑了作为 SageMaker 一部分提供的加速器。一些实例系列(如 G5 和 P4 系列)仅作为 Amazon EC2 服务的一部分提供。我们预计这些实例将在未来被 Amazon SageMaker 支持。

AWS Inferentia

AWS Inferentia 是一个专门为深度学习推理工作负载设计的 ASIC 加速器。根据 AWS 的说法,它提供了云中最低的推理成本。每个 Inferentia 芯片由四个 NeuronCore 组成,它们是高性能的矩阵乘法引擎。NeuronCore 优化了小批量操作,以确保最低的推理延迟。Inferentia 支持 FP16、BF16 和 INT8 精度类型。g4dn.xlarge 替代方案的运行成本则低 70%。

要在 Inferentia 实例上运行推理,你需要使用 AWS Neuron SDK 编译模型(github.com/aws/aws-neuron-sdk/)。Neuron SDK 支持 TensorFlow、PyTorch 和 MXNet 深度学习框架。我们将在下一节讨论使用 Neuron SDK 进行模型编译和优化。

AWS Inferentia 提供了一种高性能且具有成本效益的推理加速器。此外,你还可以使用 Neuron SDK 进一步优化模型。请注意,你需要考虑给定模型架构及其操作符是否被 Neuron SDK 支持。对于不受支持的操作符,它们将由 CPU 设备执行,这将导致额外的延迟。根据目标 SLA,可能可以接受,也可能无法接受。

亚马逊弹性推理

弹性推理 (EI) 是一种功能,允许你将用户定义的加速器能力附加到常规的 CPU 实例上。EI 专为推理用例设计。加速器能力通过附加的网络接口提供。EI 支持 TensorFlow、MXNet 和 PyTorch 框架以及 ONNX 模型格式。要使用 EI,你需要将模型加载到专门为 EI 启用的深度学习框架版本中。这些修改后的深度学习框架会自动检测到 EI 加速器的存在,并通过网络接口执行操作。下图展示了这一点:

图 8.1 – 通过网络接口访问 EI GPU 能力

图 8.1 – 通过网络接口访问 EI GPU 能力

EI 提供多种加速器类型。你可以根据所需的加速器内存或预期的吞吐量(以 TFLOPS 为单位)选择合适的类型。EI 在实例配置方面提供了低成本和高灵活性。与配置受限的专用 GPU 实例不同,你可以将 CPU 实例和 EI 混合使用,以实现可接受的推理延迟和吞吐量,同时保持整体成本较低:

加速器类型 FP32 吞吐量(TFLOPS) FP16 吞吐量(TFLOPS) 内存(GB)
eia2.medium 1 8 2
eia2.large 2 16 4
eia2.xlarge 4 32 8

图 8.2 – EI 性能特性

在选择 EI 时,你需要牢记几个注意事项:

  • 由于设计原因,EI 加速器通常会由于网络传输引入额外的延迟。对于具有复杂控制流的模型,EI 加速器可能表现不佳。

  • EI 启用的深度学习框架远远落后于最新的开源版本。此外,在 EI 上运行最新的模型架构时,你可能会遇到兼容性问题。

  • EI 提供的 GPU 内存相对较低(与最新一代 GPU 实例相比),这可能限制你在其上运行的模型类型。

与 GPU 实例和 Inferentia 一样,EI 支持模型编译和优化。你可以使用 SageMaker Neo 优化器任务来优化 TensorFlow 模型,该任务使用 TF-TRT 库进行 TensorRT 优化。优化后的模型通常具有更好的延迟-吞吐量特性,但在推理时可能会占用大量 GPU 内存,这可能会导致 内存溢出 (OOM) 问题。

当选择深度学习加速器时,EI 是一个有用的选项,特别是当你在寻找一个高度灵活且具有成本效益的解决方案,并且运行较为紧凑且需求较低的模型架构时。然而,如果你正在寻找高性能的推理处理,尤其是对高要求模型的推理,应该优先考虑 Inferentia 和 G4 实例。

为推理编译模型

为了在给定的加速器硬件上实现最佳推理性能,你通常需要为该加速器编译模型。编译过程包括各种计算优化,例如层和张量融合、精度校准,以及丢弃未使用的参数。

在本节中,我们将回顾为之前讨论的推理加速器执行编译的优化器:NVIDIA TensorRT(用于 NVIDIA GPU 加速器)和 Neuron SDK 编译器(用于 AWS Inferentia)。之后,我们将回顾一种名为 SageMaker Neo 的托管编译服务,它支持多种云和边缘硬件加速器。

我们将从查看用于 NVIDIA GPU 加速器的 TensorRT 编译器开始。

使用 TensorRT

NVIDIA TensorRT 是为 CUDA 生态系统构建的编译器和推理运行时。根据 NVIDIA 的基准测试,相比于未编译的模型版本,它能够在相同的硬件加速器上将模型性能提高最多六倍。TensorRT 支持 TensorFlow 和 PyTorch 框架,以及跨框架的 ONNX 模型格式。TensorRT 集成了 NVIDIA Triton 模型服务器,用于管理模型部署并提供推理请求服务。TensorRT 提供了 C++ 和 Python 运行时环境。C++ 运行时在边缘设备和嵌入式设备上尤其有用,这些设备可能没有配置 Python 运行时:

图 8.3 – 通过网络接口访问 EI GPU 容量

图 8.3 – 通过网络接口访问 EI GPU 容量

TensorRT 在编译模型时提供了几种关键的优化机制(参见图 8.3):

  • 精度校准将权重和激活值转换为 INT8 精度类型,而不会影响准确性,从而最大化模型吞吐量。

  • 层和张量融合将多个层和张量运算合并为单一计算,以优化内存利用率和延迟。

  • 内核自动调优为给定的硬件加速器选择最佳的数据层和算法。

  • 动态张量内存允许你高效地重用为张量分配的内存。

  • 多流执行允许你并行处理多个输入。

这些优化大多数都会在没有用户输入的情况下自动发生。在编译时,你需要设置以下参数:

  • 精度模式定义了模型参数将转换成的精度类型。TensorRT 允许你在几乎不影响准确性的情况下降低精度。较低的精度可以减少内存占用,从而加速内存绑定的操作。

  • 输入批量大小设置单次推理请求中预期的样本输入数量。增大批量大小通常会提高整体系统吞吐量。然而,较大的批量大小需要更多的可用内存,并且可能增加推理请求的延迟。

  • 最大内存大小定义了在推理时可用于模型的 GPU 内存量。

建议根据可用资源和延迟吞吐量服务水平协议(SLA),尝试各种这些参数的组合,以获得最佳性能。

根据 DL 框架模型,编译到 TensorRT 格式的路径不同。对于 TensorFlow,你可以使用 TensorFlow-TensorRTTRT)集成库(github.com/tensorflow/tensorrt)。对于 PyTorch,你需要使用 PyTorch JIT 编译器将模型转换为 TorchScript 格式。然后,你可以使用 Torch-TensorRT 集成库(github.com/pytorch/TensorRT)将模型编译为 TensorRT 格式。然后,编译后的模型可以使用你选择的模型服务器进行服务。在 第九章实现模型服务器中,我们将使用 NVIDIA Triton 模型服务器开发用于 TensorRT 编译模型的推理应用程序。

让我们回顾一下如何使用 TensorRT 编译 PyTorch ResNet50 模型的示例,并将其与未编译的模型进行基准测试。要使用 TensorRT 编译模型,你需要访问包含目标 NVIDIA GPU 的环境。以 Amazon SageMaker 为例,你可以使用具有 NVIDIA GPU 加速器的 SageMaker 笔记本实例。建议使用官方的 NVIDIA PyTorch 容器,其中预配置了所有依赖项。

重要提示

请注意,Amazon SageMaker Studio 笔记本不允许运行 Docker 容器。因此,在本示例中,我们将使用 SageMaker 笔记本实例。选择一个具有与目标推理集群相同 GPU 加速器的笔记本实例。

按照以下步骤为 TensorRT 运行时编译 PyTorch 模型:

  1. 启动一个带有 NVIDIA GPU 加速器的 SageMaker 笔记本实例。例如,你可以使用 ml.p3.2xlarge 笔记本实例。

  2. 一旦你的笔记本完全配置好,请通过 AWS 控制台中的相应链接打开 JupyterLab 服务。

  3. 在你的 JupyterLab 环境中,打开一个终端会话并运行以下命令以复制模型编译的源代码:

    cd ~/SageMaker
    git clone https://github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker
    
  4. 在同一终端会话中,运行以下命令以下载配置了 TensorRT 的 NVIDIA PyTorch 容器:

    docker pull nvcr.io/nvidia/pytorch:22.06-py3
    docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm -v ~/SageMaker/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/chapter8/1_src:/workspace/tensorrt_benchmark nvcr.io/nvidia/pytorch:22.06-py3
    
  5. 一个新的终端会话将在 PyTorch 容器中打开。运行以下命令以下载测试图像并开始基准测试:

    cd tensorrt_benchmark/
    bash data_download.sh
    python benchmarking_resnet50.py
    

基准测试脚本需要几分钟才能完成。你将能够获得未编译的 ResNet50 模型与编译后的模型(分别使用 FP32 精度和 FP16 精度)的推理结果。正如以下总结所示,与相同精度的未编译模型相比,FP16 模型的延迟提高了五倍以上:

  • 未编译的 ResNet50 模型:平均批处理时间:102.17 毫秒

  • 编译后的 ResNet50 模型(使用 FP32 精度):平均批处理时间:70.79 毫秒

  • ResNet50 模型(FP16 精度):平均批处理时间:17.26 毫秒

让我们回顾一下基准测试脚本中的编译和推理部分,熟悉 PyTorch TensorRT API:

  1. 首先,我们将从 PyTorch Hub 加载常规的、未编译的 ResNet50 模型:

    import torch
    resnet50_model = torch.hub.load("pytorch/vision:v0.10.0", "resnet50", pretrained=True)
    resnet50_model.eval()
    
  2. 要编译模型,我们可以使用torch_tensorrt集成 API。在以下示例中,我们将模型编译成一个 TorchScript 模块,以便针对 TensorRT 引擎进行优化:

    import torch_tensorrt
    trt_model_fp32 = torch_tensorrt.compile(model, inputs = [torch_tensorrt.Input((128, 3, 224, 224), dtype=torch.float32)],
        enabled_precisions = torch.float32,
        workspace_size = 1 << 22
    )
    
  3. 现在,您可以像普通 TorchScript 程序一样保存并加载编译后的模型:

    trt_model_fp32.save('resnet50_fp32.pt')
    loaded = torch.jit.load('resnet50_fp32.pt')
    

在本节中,您了解了如何使用 TensorRT 手动编译 PyTorch 模型以适配 NVIDIA GPU 加速器,并且回顾了编译后模型的延迟改进。

如果您有兴趣编译 TensorFlow 模型,您可以使用类似的方法。请注意,您需要使用官方的 NVIDIA TensorFlow 容器。关于此的代码示例,您可以参考官方的 TensorFlow 教程:blog.tensorflow.org/2021/01/leveraging-tensorflow-tensorrt-integration.xhtml

正如您所看到的,整个编译过程是手动的。本章后面我们将介绍 SageMaker Neo,它可以让我们以最少的手动操作编译 TensorFlow 和 PyTorch 模型以适配 NVIDIA GPU 加速器。

使用 Neuron SDK

AWS Neuron SDK 允许您将 DL 模型编译为 AWS Inferentia 实例。它提供了多个参数,帮助您根据可用的 Inferentia 芯片以及您的延迟和吞吐量 SLA 来优化推理程序。Neuron SDK 支持 TensorFlow、PyTorch 和 MXNet 框架。Neuron SDK 是一个提前编译的工具,因此您必须在编译时显式提供批量大小。它还包括一个运行时环境,我们在其中加载模型并在推理时获取预测。请注意,Neuron SDK 编译的模型只能在 AWS Inferentia 芯片上使用。

Neuron SDK 支持一组广泛但有限的操作符。AWS 在以下流行的模型架构上测试了 Neuron SDK:

  • 来自 HuggingFace Transformer 库的 NLP 模型BERTdistilBERTXLM-BERTRobertBioBERTMarianMTPegasus和 Bart

  • 计算机视觉模型ResnetRenextVGGYolo v3/v4/v5SSD

Neuron SDK 还支持通用模型层,如全连接层或嵌入查找。如果您的模型架构使用了支持的操作符,您将能够充分利用 Neuron SDK 的优化。您可以参考 Neuron SDK 官方文档中的支持操作符列表,查看具体的 DL 框架:awsdocs-neuron.readthedocs-hosted.com/en/latest/index.xhtml

在使用 Neuron SDK 编译模型时,请牢记以下几点注意事项:

  • 如果某个特定操作符不被支持,则该操作的执行将转移到 CPU 加速器上,这会导致性能变慢。

  • 您模型中的控制流可能无法完全得到支持。

  • 如果您期望变量批量大小,您需要实现动态批处理

  • 如果您期望输入大小可变(例如,输入图像的大小不固定),您应该考虑实现填充或分桶。

现在,让我们讨论可用的 Neuron SDK 优化。

FP32 自动类型转换

每当可能时,Neuron SDK 将您的模型转换为 BF16 精度类型,以减少内存占用并改善延迟-吞吐特性。

批量推理输入

批处理是指将多个推理输入合并为一个批次。在这方面,它与模型训练中的批处理相同。对于推理工作负载,批处理会影响您的吞吐量。像 TensorRT 一样,Neuron SDK 要求您在编译时定义目标批量大小。Inferentia 加速器特别优化了对较小批量大小的推理运行。这是通过将延迟敏感的操作(如从内存中读取权重)组合到整个推理批次中,从而比对每个推理输入执行相同操作时获得更好的延迟-吞吐特性。下图说明了这一概念:

图 8.4 – 使用单一内存检索的批量推理

图 8.4 – 使用单一内存检索的批量推理

动态批处理是 Neuron SDK 的一项功能,它允许您切分输入张量,使其匹配编译时使用的批量大小。请注意,动态批处理适用于若干合适的模型架构。

NeuronCore 流水线

每个 Inferentia 加速器由四个NeuronCores组成。流水线技术允许您将模型分片到多个 NeuronCore 上,将模型参数缓存到片上内存中。这使得您可以使用本地缓存的数据更快地处理网络运算符,并避免访问外部内存。根据 AWS 的说法,内部基准流水线通常使我们在没有批处理的情况下实现最高的硬件利用率。下图展示了流水线的示例:

图 8.5 – 将模型流水线化到三个 NeuronCores 中

图 8.5 – 将模型流水线化到三个 NeuronCores 中

在以下示例中,我们将在 AWS Inferentia 实例上编译并基准测试 ResNet50 模型。撰写本文时,Amazon SageMaker 不支持托管的笔记本实例。因此,我们使用了 Amazon EC2 inf1.xlarge实例,配置为8888。为此,您需要像这样设置实例的安全组:

图 8.6 – 配置安全组以允许 Jupyter 流量

图 8.6 – 配置安全组以允许 Jupyter 流量

在开始编译 Neuron SDK 之前,我们需要在 EC2 实例上安装 Neuron SDK 及其依赖项。按照以下步骤进行操作:

  1. 首先,你需要使用以下命令 SSH 连接到你的实例:

    chmod 400 <your_ssh_key>
    ssh -i <your_ssh_key>ubuntu@<your_instance_public_DNS>
    

登录到 EC2 实例后,请按照 awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-intro/pytorch-setup/pytorch-install.xhtml 上的说明,在你的 Ubuntu 操作系统上安装 Neuron PyTorch。请注意,安装过程可能需要大约 5 分钟完成。

  1. 安装完成后,克隆源代码并启动 Jupyter 服务器应用程序:

    git clone https://github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker
    cd Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/chapter8/
    jupyter notebook --ip=0.0.0.0
    
  2. 之后,你可以打开 <your_instance_public_DNS>:8888/tree 以访问此示例的 Jupyter notebook。请注意,第一次进行此操作时,你需要复制之前由 jupyter notebook... 返回的安全令牌。

设置完成后,我们可以在 AWS Inferentia 加速器上编译并基准测试模型。完整代码可在此查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter8/2_Neuron_SDK_compilation.ipynb。按照以下步骤操作:

  1. 在打开的 Jupyter notebook 中,将内核更改为我们之前配置的Python (Neuron PyTorch)

  2. 接下来,我们必须导入所需的库,包括 torch_neuron,并下载 ResNet50 模型:

    import torch
    from torchvision import models, transforms, datasets
    import torch_neuron
    image = torch.zeros([1, 3, 224, 224], dtype=torch.float32)
    model = models.resnet50(pretrained=True)
    model.eval()
    
  3. 然后,我们必须分析模型操作符,确定是否有任何操作符不被 Inferentia/Neuron SDK 支持。由于 ResNet50 模型已被支持,此命令的输出应确认所有模型操作符均受支持:

    torch.neuron.analyze_model(model, example_inputs=[image])
    
  4. 现在,我们已准备好通过运行以下命令来进行编译。你将看到编译统计信息和状态输出:

    model_neuron = torch.neuron.trace(model, example_inputs=[image])
    
  5. 由于 Neuron SDK 编译成了一个 TorchScript 程序,保存和加载模型的方式类似于你在常规 PyTorch 中的操作:

    model_neuron.save("resnet50_neuron.pt")
    model_neuron = torch.jit.load('resnet50_neuron.pt')
    
  6. 现在,让我们通过批处理或流水线对已编译的模型进行基准测试。为此,我们将准备预处理和基准测试方法,形成推理批处理,并测量延迟ms)和吞吐量samples/s):

    model_neuron_parallel = torch.neuron.DataParallel(model_neuron)
    num_neuron_cores = 4
    image = preprocess(batch_size=batch_size, num_neuron_cores=num_neuron_cores)
    benchmark(model_neuron_parallel, image)
    

基准测试结果应类似于以下内容:

Input image shape is [4, 3, 224, 224]
Avg. Throughput: 551, Max Throughput: 562
Latency P50: 7
Latency P90: 7
Latency P95: 7
Latency P99: 7
  1. 接下来,我们将重新编译模型,并启用批处理功能(通过将 batch_size 设置为每个 NeuronCore 5 个样本):

    batch_size = 5
    image = torch.zeros([batch_size, 3, 224, 224], dtype=torch.float32)
    model_neuron = torch.neuron.trace(model, example_inputs=[image])
    model_neuron.save("resnet50_neuron_b{}.pt".format(batch_size))
    
  2. 在重新运行基准测试后,请注意,尽管延迟有所减少,但整体吞吐量已增加,如下所示:

    Batch_size = 5
    model_neuron = torch.jit.load("resnet50_neuron_b{}.pt".format(batch_size))
    model_neuron_parallel = torch.neuron.DataParallel(model_neuron)
    image = preprocess(batch_size=batch_size, num_neuron_cores=num_neuron_cores)
    benchmark(model_neuron_parallel, image)
    

基准测试的输出应如下所示:

Input image shape is [20, 3, 224, 224]
Avg. Throughput: 979, Max Throughput: 998
Latency P50: 20
Latency P90: 21
Latency P95: 21
Latency P99: 24
  1. 最后,让我们在启用流水线的情况下编译并基准测试模型。我们将首先通过 neuroncore-pipeline-cores 参数追踪原始模型:

    neuron_pipeline_model = torch.neuron.trace(model,
                                               example_inputs=[image],
                                               verbose=1,
                                               compiler_args = ['--neuroncore-pipeline-cores', str(num_neuron_cores)])
    
  2. 然后,我们将在这个新模型上重新运行基准测试:

    image = preprocess(batch_size=batch_size, num_neuron_cores=num_neuron_cores)
    benchmark(neuron_pipeline_model, image)
    

此基准测试的输出将如下所示:

Input image shape is [20, 3, 224, 224]
Avg. Throughput: 271, Max Throughput: 274
Latency P50: 73
Latency P90: 74
Latency P95: 74
Latency P99: 79

请注意,管道模型的最终延迟和吞吐量低于无批处理和有批处理的模型。出现这种情况的一个原因是,在我们的基准测试中,推理请求是按顺序运行的。为了更好地利用管道模型,我们需要创建多个并行推理请求。

使用 SageMaker Neo

SageMaker Neo 允许你编译和优化深度学习模型,以适配各种硬件平台。它支持 PyTorch、TensorFlow、MXNet 和 ONNX 模型,适用于 Ambarella、ARM、Intel、NVIDIA、NXP、Qualcomm、Texas Instruments 和 Xilinx 等硬件平台。SageMaker Neo 还支持云实例部署以及边缘设备部署。

在底层,SageMaker Neo 将训练好的模型从框架特定的表示转换为中间的框架无关表示。接着,它会应用自动优化,并生成优化操作的二进制代码。模型编译完成后,你可以通过 SageMaker 推理服务将其部署到目标实例类型。Neo 还为每个目标平台提供运行时,加载并执行编译后的模型。下图展示了 SageMaker Neo 的概览:

图 8.7 – SageMaker Neo 概览

图 8.7 – SageMaker Neo 概览

SageMaker Neo 能够显著减少额外的开发或设置工作。然而,它也有一些限制,这些限制可能适合或不适合你的具体使用案例。

  • SageMaker Neo 主要支持计算机视觉模型,如图像分类目标检测语义分割。它不支持 NLP 模型架构等其他类型的模型。

  • SageMaker Neo 支持多个深度学习框架,但它们通常落后几个主要版本。因此,如果你希望使用最新的模型架构和/或框架版本特性,你需要考虑其他编译选项(例如,使用 TensorRT 手动编译)。有关 SageMaker Neo 支持的最新详情,请参考 docs.aws.amazon.com/sagemaker/latest/dg/neo-supported-cloud.xhtml

  • SageMaker Neo 支持一系列云实例。截至目前,它支持 ml.c5ml.c4ml.m5ml.m4ml.p3ml.p2ml.inf1 实例的编译。

  • SageMaker Neo 对推理请求格式有特定要求(特别是在输入形状方面)。

在考虑使用 SageMaker Neo 时,需牢记以下关键限制。在很多情况下,如果你的模型架构、框架版本和目标硬件加速器受到支持,SageMaker Neo 可以是一个方便高效的编译方式。

让我们回顾一下如何使用 SageMaker Neo 编译 TensorFlow 模型。在这个例子中,我们将训练 ResNet50 模型,针对多个硬件平台进行编译,并部署优化后的模型推理端点。我们将重点介绍关键部分。完整的源代码请参见这里:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter8/3_SageMaker_Neo_TF.ipynb

开发训练和推理脚本

在前面的章节中,我们提到过,SageMaker 必须实现特定的方法才能在 SageMaker 上训练模型和运行推理。根据深度学习框架和目标硬件平台的不同,所需的 API 会略有不同。详细信息请参见官方文档:docs.aws.amazon.com/sagemaker/latest/dg/neo-deployment-hosting-services-prerequisites.xhtml

为了服务 TensorFlow 模型,我们实现了简单的serving_input_fn()方法,该方法将输入传递给模型并返回预测结果:

def serving_input_fn():
    inputs = {"x": tf.placeholder(tf.float32, [None, 784])}
    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

接下来,我们安排编译作业。

运行编译作业和部署

要开始编译作业,我们必须使用 SageMaker Python SDK 中的sagemaker.tensorflow来训练模型,然后导入 TensorFlow:

mnist_estimator = TensorFlow(entry_point='mnist.py',
                             source_dir="3_src",
                             role=role,
                             instance_count=1,
                             instance_type='ml.p3.2xlarge',
                             framework_version='1.15.0',
                             py_version='py3',
                             )
mnist_estimator.fit(training_data_uri)

一旦模型训练完成,我们可以测试在两种不同硬件加速器上编译的效果:一个是配备 NVIDIA GPU 设备的p2实例,另一个是没有任何专用硬件的c5实例。请按照以下步骤操作:

  1. 为此,首先,我们必须为 NVIDIA GPU 编译模型,并使用相同的硬件类型部署端点。请注意input_shape参数,它告诉 SageMaker 在编译过程中使用何种输入形状。你需要在推理时将推理样本转换为相同的输入形状:

    p2_estimator = mnist_estimator.compile_model(target_instance_family='ml_p2', 
                                  input_shape={'data':[1, 784]},
                                  output_path=output_path)
    p2_predictor = p2_estimator.deploy(initial_instance _count = 1,
                                                     instance_type = 'ml.inf1.xlarge')
    
  2. 若要访问编译作业的日志,你可以在 AWS 控制台中导航到SageMaker | 推理 | 编译作业。在这些日志中,你可以找到例如 SageMaker Neo 使用的编译框架(Apache TVM),并查看模型操作符的编译状态。

  3. 对于 c5 实例,运行编译作业非常相似。请注意,我们使用的是与编译 p2 实例时相同的estimator对象。如前所述,你只需要训练一次模型;然后,你可以为多个目标平台进行编译:

    c5_estimator = mnist_estimator.compile_model(target_instance_family='ml_c5', 
                                  input_shape={'data':[1, 784]},  
                                  output_path=output_path)
    c5_predictor = c5_estimator.deploy(initial_instance_count = 1,
                                                     instance _type = 'ml.c5.xlarge')
    
  4. 成功编译后,生成的模型工件将保存在调用c5_estimator.model_data属性时可访问的 S3 位置。

  5. 调用已编译模型的端点与调用未编译模型的端点相同。以下是p2推理端点的示例:

    data = inference_data[i].reshape(1,784)
    predict_response= p2_predictor.predict(data)
    

请注意,我们重新调整了输入数据,以便与编译过程中使用的输入数据匹配。

这个简短的例子演示了如何使用 SageMaker Neo 编译单个模型。建议在将模型部署到生产环境之前,对 SageMaker 编译的模型进行基准测试,以确认延迟吞吐量的改进。请注意,未编译和已编译的模型可能具有不同的内存需求。

概要

在本章中,我们回顾了适用于运行 DL 推断程序的可用硬件加速器。我们还讨论了如何使用 TensorRT 编译器优化您的模型,以适配 NVIDIA GPU 加速器和 Neuron SDK 适配 AWS Inferentia 加速器。然后,我们回顾了 SageMaker Neo 服务,该服务允许您为广泛的硬件平台编译支持的模型,减少开发工作,并突出了该服务的几个限制。阅读完本章后,您应该能够根据您特定的使用案例需求(如延迟、吞吐量和成本)来做出关于使用哪种硬件加速器以及如何优化它们的决策。

一旦选择了您的硬件加速器和模型优化策略,您将需要决定使用哪种模型服务器以及如何在服务时间进一步调整您的推理工作负载。在下一章中,我们将讨论流行的模型服务器解决方案,并获得在 SageMaker 推理服务上开发和部署它们的实际经验。

第九章:实现模型服务器

第八章,《考虑推理硬件》中,我们讨论了为服务深度学习(DL)模型提供的硬件选项和优化,这些选项作为 Amazon SageMaker 平台的一部分供您使用。在本章中,我们将关注工程推理工作负载的另一个重要方面——选择和配置模型服务器。

模型服务器类似于常规应用程序的应用服务器,为您的 DL 模型提供运行时环境。作为开发者,您将训练好的模型部署到模型服务器,后者将已部署的模型暴露为 REST 或 gRPC 端点。DL 模型的最终用户随后会向这些已建立的端点发送推理请求,并收到包含预测结果的响应。模型服务器可以同时为多个最终用户提供服务。它还提供可配置的机制,以优化推理延迟和吞吐量,满足特定的服务水平协议(SLA)。

第一章,《介绍 Amazon SageMaker 深度学习》中,我们讨论了 Amazon SageMaker 托管服务的几种模型部署机制:实时推理端点、批量转换作业和异步推理。在所有这些情况下,您都需要选择一个模型服务器来管理推理运行时和模型部署。然而,这些使用案例的模型服务器配置可能会有所不同,因为它们具有不同的推理流量特征和延迟/吞吐量要求。

Amazon SageMaker 提供了多种模型服务器解决方案,作为其深度学习推理容器的一部分。在本章中,我们将重点介绍三种流行的模型服务器,这些服务器旨在将深度学习推理工作负载投入生产:TensorFlow Serving (TFS)、PyTorch TorchServe (PTS) 和 NVIDIA Triton

在本章中,我们将讨论以下主题:

  • 使用 TFS

  • 使用 PTS

  • 使用 NVIDIA Triton

阅读完本章后,您将了解如何部署您的 TensorFlow 和 PyTorch 模型,并根据您的推理需求配置模型服务器。我们还将讨论将模型服务器作为 SageMaker 托管服务一部分时的功能限制。

技术要求

在本章中,我们将提供代码示例,帮助您培养实践技能。完整的代码示例可在此处查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter9/

为了跟随本章的代码,您需要具备以下内容:

  • 拥有一个 AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户。

  • 已建立 SageMaker Notebook、SageMaker Studio Notebook 或本地兼容 SageMaker 的环境。

  • 访问你 AWS 账户中的 GPU 训练实例。本章中的每个示例将提供推荐的实例类型供你使用。你可能需要增加 SageMaker 训练任务 的计算配额,以启用 GPU 实例。在这种情况下,请按照 docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.xhtml 上的说明操作。

  • 你必须通过运行 pip install -r requirements.txt 来安装所需的 Python 库。包含所需库的文件可以在 chapter9 目录中找到。

  • 在本章中,我们将提供编译模型进行推理的示例,这需要访问特定的加速器类型。请在模型服务器示例中查看实例推荐。

使用 TFS

TFS 是一个原生的模型服务器,支持 TensorFlow 1、TensorFlow 2 和 Keras 模型。它旨在提供一个灵活且高性能的运行时环境,配备了广泛的管理 API 和操作功能(例如日志记录和度量)。AWS 提供了 TFS 作为 TensorFlow 推理容器的一部分 (github.com/aws/deep-learning-containers/tree/master/tensorflow/inference/docker)。

回顾 TFS 概念

TFS 有一个叫做 servable 的概念,它封装了推理所需的所有模型和代码资产。为了将 servable 准备好供 TFS 服务使用,你需要将训练好的模型打包成 SavedModel 格式。SavedModel 包含了一个完整的 TensorFlow 程序,包括训练过的参数和计算逻辑。它不需要原始的模型构建代码就能运行,这使得它在 TFS 生态系统中共享或部署变得十分方便(例如,使用 TFLite、TensorFlow.js 或 TFS)。你也可以在一个 servable 中打包多个模型以及特定的模型查找或嵌入。

TFS 通过 REST 或 gRPC 端点加载并暴露你的 servable。服务器 API 定义了一系列端点,用于执行分类和回归推理。此外,每个 servable 都有一个关联的 signature,它定义了模型的输入和输出张量,以及模型类型(回归或分类)。许多常见模型都有标准的 signature,这些 signature 取决于任务类型(例如,图像分类、物体检测、文本分类等)。TFS 还允许你拥有自定义的 signature。

将 TFS 与 SageMaker 集成

Amazon SageMaker 提供了一个托管的环境,在这个环境中,无论底层的模型服务器如何,你都可以通过统一的管理和调用 API 来管理推理端点。这种方法对原生模型服务器功能设定了一些限制。在本节中,我们将回顾 SageMaker 如何与 TFS 集成以及你应了解的限制。

当在 SageMaker 上部署 TFS 时,您将无法访问 TFS 的本地管理 API 来管理您的可服务对象生命周期(加载和卸载模型、推动模型版本等)。此外,您也无法直接访问 TFS 服务 API。相反,您需要使用标准的 SageMaker 调用接口来调用您的 SageMaker 端点。然后,SageMaker HTTP 服务器(DL TFS 容器的一部分)将您的请求转换为 TFS 格式,并将其传递给 TFS 服务 API。请注意,您可以在推理脚本中提供自定义的预处理、预测和后处理逻辑。SageMaker 支持 REST 和 gRPC 服务 API。下图展示了这种 TFS 集成:

图 9.1 – TFS 与 SageMaker 托管集成

图 9.1 – TFS 与 SageMaker 托管集成

在 SageMaker 上使用 TFS 时,有几件事需要注意:

  • 如前所述,SageMaker 不允许您访问 TFS 管理 API。然而,它允许您通过环境变量提供 TFS 的配置。

  • SageMaker 支持使用 TFS 托管多个模型。为此,您需要为每个模型准备单独的可服务对象,然后创建一个多模型归档。

  • 您可以使用 REST 头部和请求体来指定 TFS 应使用哪些模型来处理特定请求。例如,以下请求告诉 TFS 使用 model2 来处理此请求:

    aws sagemaker-runtime invoke-endpoint \
        --endpoint-name my-endpoint \
        --content-type 'application/json' \
        --body '{"instances": [1.0, 2.0, 5.0]}' \
        --custom-attributes 'tfs-model-name=other_model
    

SageMaker 支持默认的 TFS 输入和输出格式来处理您的推理请求。此外,SageMaker 还支持 application/JSON、text/CSV 和 application/JSON Lines 格式。

请注意,一旦您使用 TFS 模型服务器部署端点后,您将无法直接更改 TFS 配置或服务的模型。为此,您需要使用 SageMaker 管理 API 来创建一个新端点或端点变体,并使用所需的配置。我们将在第十章推理工作负载的操作化中讨论如何在生产中管理 SageMaker 推理资源。

优化 TFS

TFS 提供了一系列机制,以根据您的需求、运行时环境和可用硬件资源优化模型服务。这意味着 TFS 调优是特定用例的,通常需要测试和基准测试才能实现预期的性能。在本节中,我们将回顾几种可以用来调优 TFS 性能的机制。

使用 TFS 批处理

TFS 支持自动批处理,您可以将多个推理请求放在一个批次中。这可以提高您的服务器吞吐量,尤其是在使用 GPU 实例时(请记住,GPU 非常适合并行计算)。如何配置批处理会根据硬件设备的类型有所不同。TFS 为不同的可服务对象支持不同的批处理调度。

要在 SageMaker 上配置 TFS 批处理,您可以使用以下环境变量:

  • SAGEMAKER_TFS_ENABLE_BATCHING 用于启用 TFS 批处理功能。默认值为 false,意味着批处理未启用。

  • SAGEMAKER_TFS_MAX_BATCH_SIZE 定义了批次的最大大小。默认值为 8

  • SAGEMAKER_TFS_BATCH_TIMEOUT_MICROS 定义了在微秒内等待以积累完整批次的时间。默认值为 1000

  • SAGEMAKER_TFS_NUM_BATCH_THREADS 设置了同时处理多少个批次。默认值为实例 CPU 的数量。

  • SAGEMAKER_TFS_MAX_ENQUEUED_BATCHES 定义了同一时间可以排队的批次数量。

你可以在这里查看 TFS 批处理功能的详细文档:github.com/tensorflow/serving/blob/master/tensorflow_serving/batching/README.md

使用 gRPC 服务 API

如前所述,TFS 支持两种类型的 API:REST 和 gRPC。虽然这两种 API 功能相同,但由于使用了 HTTP/2 网络协议和通过 ProtoBuf 格式进行更高效的负载表示,gRPC API 通常具有更好的性能。

虽然 SageMaker 调用 API 仅支持 REST API,但你仍然可以使用 gRPC 进行 SageMaker HTTP 前端服务器与 TFS 之间的容器间通信(有关此内容的示意图,请参见 图 9.1)。请注意,在这种情况下,你需要提供一些代码将 SageMaker 负载转换为 gRPC 格式并将其发送到 TFS。然而,即便如此,AWS 报告称,在图像分类任务中,整体延迟至少降低了 75%。详情请参考本文:aws.amazon.com/blogs/machine-learning/reduce-compuer-vision-inference-latency-using-grpc-with-tensorflow-serving-on-amazon-sagemaker/。性能提升会根据模型和负载大小有所不同。

使用 TFS 配置资源利用率

TFS 提供以下参数来配置硬件资源分配:

  • SAGEMAKER_TFS_INSTANCE_COUNT 定义了将生成多少个 TFS 服务进程实例。改变此参数可能会增加 CPU 和 GPU 利用率,并最终改善延迟/吞吐量特性。

  • SAGEMAKER_TFS_FRACTIONAL_GPU_MEM_MARGIN 定义了可用于初始化 CUDA/cuDNN 库的 GPU 内存的比例。剩余内存将平均分配给 TFS 进程。

  • SAGEMAKER_TFS_INTER_OP_PARALLELISM 决定在运行模型图中的独立非阻塞计算操作时使用多少线程。

  • SAGEMAKER_TFS_INTRA_OP_PARALLELISM 决定在运行可以内部并行化的操作时使用多少线程。

现在,让我们通过一个实际例子来回顾如何在 SageMaker 上使用 TFS。

实现 TFS 服务

在本示例中,我们将从 TensorFlow Hub 获取一个预训练的模型,将其转换为 SavedModel 格式,然后与自定义推理一起打包,部署到 SageMaker 上。我们将回顾如何使用 REST 和 gRPC API,以及如何在部署到 SageMaker 管理型托管时定义 TFS 配置。为了完成这项任务,我们将使用流行的 EfficientNetV2 模型架构来进行图像分类。

完整代码可以在这里找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter9/1_TensorFlow_Serving.ipynb

准备训练模型

我们将从 TensorFlow Hub 加载模型工件。您可以在其模型页面阅读有关 EfficientNetV2 模型的信息,地址是:tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/classification/2。为了下载模型,我们可以使用 TensorFlow Hub API,如下所示的代码块所示:

import tensorflow as tf
import tensorflow_hub as hub
model_handle = "https://tfhub.dev/google/imagenet/efficientnet_v2_imagenet1k_s/classification/2"
classifier = hub.load(model_handle)

该模型期望一个密集的 4D 张量,数据类型为 float32,形状为 [batch, height, weight, color],其中 heightweight 的固定长度为 384color 的长度为 3batch 可以是可变的。

为了在本地测试模型,您需要将图像(或一批图像)转换为期望的 4D 张量,将其通过模型运行,并应用 softmax 函数以获取标签概率,如下所示:

probabilities = tf.nn.softmax(classifier(image)).numpy()

现在我们已经对模型进行了基本的烟雾测试,接下来需要将其打包为 SageMaker/TFS 兼容的格式。

打包模型工件

如前所述,TFS 期望您的模型被转换为 SavedModel 格式。此外,SageMaker 期望模型工件被打包成一个 tar.gz 压缩包,并具有以下结构:

    model1
        |--[model_version_number]
            |--variables
            |--saved_model.pb
    model2
        |--[model_version_number]
            |--assets
            |--variables
            |--saved_model.pb
    code
        |--inference.py
        |--requirements.txt

以下代码创建了适当的目录结构并将训练好的模型导出为 SavedModel 格式:

model_name = "efficientnetv2-s"
model_dir = f"./{model_name}/1"
code_dir = f"./{model_name}/code"
os.makedirs(model_dir, exist_ok=False)
os.makedirs(code_dir, exist_ok=False)
tf.saved_model.save(classifier, model_dir)

请注意,在我们的示例中,我们将只使用单个版本的单个模型。接下来,我们需要为预处理、运行预测和后处理准备一个推理脚本,用于在 SageMaker HTTP 前端与 TFS 服务器之间进行交互。

开发推理代码

SageMaker 期望您的处理代码被命名为 inference.py 并放置在模型归档中的 /code 目录下。我们的推理代码需要实现 input_handler()output_handler() 函数,或一个单一的 handler() 函数。在我们的例子中,我们选择实现一个单一的 handler() 方法来处理传入的请求并将其发送到适当的 TFS API:

def handler(data, context):
    if context.request_content_type == "application/json":
        instance = json.loads(data.read().decode("utf-8"))
    else:
        raise ValueError(
            415,
            'Unsupported content type "{}"'.format(
                context.request_content_type or "Unknown"
            ),
        )
    if USE_GRPC:
        prediction = _predict_using_grpc(context, instance)
    else:
        inst_json = json.dumps({"instances": instance})
        response = requests.post(context.rest_uri, data=inst_json)
        if response.status_code != 200:
            raise Exception(response.content.decode("utf-8"))
        prediction = response.content
    response_content_type = context.accept_header
    return prediction, response_content_type

如你所见,依据我们选择使用 gRPC API 还是 REST API,处理和预测代码会略有不同。请注意,context namedtuple 对象提供了有关 TFS 配置的必要细节,例如端点路径和端口、模型名称和版本等。

如果选择使用 TFS REST API,我们需要将传入的请求转换为预期的 TFS 格式,将其序列化为 JSON,然后生成一个 POST 请求。

若要使用 gRPC API,我们需要将传入的 REST 负载转换为protobuf对象。为此,我们将使用以下辅助函数:

from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
def _predict_using_grpc(context, instance):
    grpc_request = predict_pb2.PredictRequest()
    grpc_request.model_spec.name = "model"
    grpc_request.model_spec.signature_name = "serving_default"
    options = [
        ("grpc.max_send_message_length", MAX_GRPC_MESSAGE_LENGTH),
        ("grpc.max_receive_message_length", MAX_GRPC_MESSAGE_LENGTH),
    ]
    channel = grpc.insecure_channel(f"0.0.0.0:{context.grpc_port}", options=options)
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
grpc_request.inputs["input_1"].CopyFrom(tf.make_tensor_proto(instance))
    result = stub.Predict(grpc_request, 10)
    output_shape = [dim.size for dim in result.outputs["output_1"].tensor_shape.dim]
    np_result = np.array(result.outputs["output_1"].float_val).reshape(output_shape)
    return json.dumps({"predictions": np_result.tolist()})

在这里,我们使用prediction_service_pb2()predict_pb2() TFS 方法与 gRPC API 进行通信。stub对象在 RPC 过程中转换参数。grpc_request对象定义了要调用的 TFS API 以及传递的参数。

为了选择要调用的 TFS API,我们实现了一个简单的机制,允许你通过 SageMaker 模型对象提供USE_GRPC环境变量:

USE_GRPC = True if os.getenv("USE_GRPC").lower() == "true" else False 

一旦我们的inference.py代码准备好后,我们可以将其添加到模型包中,并创建一个tar.gz模型归档。可以通过在 Jupyter notebook 中运行以下 Bash 代码来完成此操作:

! cp 1_src/inference.py $code_dir
! cp 1_src/requirements.txt $code_dir
! tar -C "$PWD" -czf model.tar.gz  efficientnetv2-s/

现在,我们的模型已经按照 TFS 和 SageMaker 的要求打包完毕,准备进行部署。

部署 TFS 模型

部署 TFS 模型,请按照以下步骤操作:

  1. 我们将首先把模型归档上传到 Amazon S3,以便 SageMaker 在部署时将其下载到服务容器中。我们可以使用 SageMaker 的Session()对象来执行此操作:

    import sagemaker
    from sagemaker import get_execution_role
    sagemaker_session = sagemaker.Session()
    role = get_execution_role() 
    bucket = sagemaker_session.default_bucket()
    prefix = 'tf-serving'
    s3_path = 's3://{}/{}'.format(bucket, prefix)
    model_data = sagemaker_session.upload_data('model.tar.gz',
                                               bucket,
                                               os.path.join(prefix, 'model'))
    
  2. 然后,我们可以使用 SageMaker SDK 的 TensorFlowModel 对象来配置 TFS 环境。请注意,我们通过env字典提供 TFS 配置:

    from sagemaker.tensorflow import TensorFlowModel
    env = {
           "SAGEMAKER_TFS_ENABLE_BATCHING":"true",
           "SAGEMAKER_TFS_MAX_BATCH_SIZE":"4",
           "SAGEMAKER_TFS_BATCH_TIMEOUT_MICROS":"100000",
           "SAGEMAKER_TFS_NUM_BATCH_THREADS":"6",
           "SAGEMAKER_TFS_MAX_ENQUEUED_BATCHES":"6",
           "USE_GRPC":"true" # to switch between TFS REST and gRCP API
           }
    tensorflow_serving_model = TensorFlowModel(model_data=model_data,
                                     name="efficientnetv2-1",
                                     role=role,
                                     framework_version='2.8',
                                     env=env,
    sagemaker_session=sagemaker_session)
    

配置好模型后,我们可以准备部署端点。在这里,我们将使用一个 GPU 实例,但你也可以尝试使用 CPU 实例。

在我们能够运行预测之前,我们需要将图像(或多个图像)转换为 4D TFS 张量,并将其转换为 NumPy ndarray,这样.predict()方法才能将其序列化为 application/JSON 内容类型。样本笔记本中提供了一个将图像处理为 TFS 格式的方法示例。

在以下代码中,我们正在运行预测,并将结果的 softmax 分数映射到标签:

response_remote = predictor.predict(image.numpy())
probabilities = np.array(response_remote['predictions'])
top_5 = tf.argsort(probabilities, axis=-1, direction="DESCENDING")[0][:5].numpy()
np_classes = np.array(classes)
# Some models include an additional 'background' class in the predictions, so
# we must account for this when reading the class labels.
includes_background_class = probabilities.shape[1] == 1001
for i, item in enumerate(top_5):
  class_index = item if includes_background_class else item + 1
  line = f'({i+1}) {class_index:4} - {classes[class_index]}: {probabilities[0][top_5][i]}'
  print(line)

运行此代码后,你应该得到包含标签及其归一化概率的输出。

在本节中,我们回顾了如何在 Amazon SageMaker 上使用 TFS 模型服务器。 TFS 是一个高度可配置的生产级模型服务器,当涉及到托管 TensorFlow 模型时,它应该被视为一个很好的选择。 我们还讨论了一些 Sagemaker / TFS 集成的实现细节,这些细节在设计您的模型服务器时应予以考虑。 一旦您在 SageMaker 上运行您的 TensorFlow 模型,建议进行基准测试,并根据您特定的用例需求调整 TFS 配置。

在接下来的部分中,我们将审查用于 PyTorch 模型的本地模型服务器 - TorchServe。

使用 PTS

PTS是 PyTorch 模型的本地模型服务器。 PTS 是 Meta 和 AWS 合作开发的,旨在为 PyTorch 生态系统提供一个生产就绪的模型服务器。 它允许您通过 REST 或 gRPC 端点提供服务和管理多个模型。 PTS 支持为了更好的推断性能提供 TorchScripted 模型。 它还配备了收集日志和指标以及优化调整的实用程序。 SageMaker 将 PTS 作为 PyTorch 推断容器的一部分支持(github.com/aws/deep-learning-containers/tree/master/pytorch/inference/docker)。

与 SageMaker 集成

PTS 是 Amazon SageMaker 上 PyTorch 模型的默认模型服务器。 与 TFS 类似,SageMaker 不向最终用户公开原生 PTS API 以进行模型管理和推断。 下图显示了如何集成 SageMaker 和 PTS:

图 9.2 - SageMaker 上的 PTS 架构

图 9.2 - SageMaker 上的 PTS 架构

让我们突出这些集成细节:

  • SageMaker 原生支持有限数量的 PTS 配置。 如果您需要对 PTS 配置更具灵活性,可能需要扩展 SageMaker PyTorch 推断容器。 或者,您可以将 PTS 配置作为模型包的一部分打包,并通过TS_CONFIG_FILE环境变量提供路径。 但是,通过后一种方法,您将无法操纵所有设置(例如,JVM 配置)。

  • PTS 要求您将模型工件和处理程序代码打包到 MAR 存档中。 SageMaker 对于模型存档有略有不同的要求,我们将在以下代码示例中讨论。

  • SageMaker 支持同时托管多个模型。 为此,您需要将ENABLE_MULTI_MODEL环境变量设置为true,并将您的模型打包到单个存档中。

SageMaker 通过端点环境变量提供配置 PTS 的机制。 让我们回顾一下可用的配置参数。

在 SageMaker 上优化 PTS

PTS 支持两种主要的性能优化机制:服务器端批处理和生成多个模型线程。 可以通过以下环境变量配置这些设置:

  • SAGEMAKER_TS_BATCH_SIZE 设置服务器端批次的最大大小。

  • SAGEMAKER_TS_MAX_BATCH_DELAY 设置服务器等待完成批次的最大延迟(以微秒为单位)。

  • SAGEMAKER_TS_RESPONSE_TIMEOUT 设置推理响应不可用时的超时时间延迟(以秒为单位)。

  • SAGEMAKER_TS_MIN_WORKERSSAGEMAKER_TS_MAX_WORKERS 分别配置 CPU 或 GPU 设备上模型工作线程的最小和最大数量。你可以在 PyTorch 文档中阅读有关如何设置这些参数的注意事项,链接如下:github.com/pytorch/serve/blob/master/docs/performance_guide.md

此外,PTS 支持使用 PyTorch TensorBoard 插件进行推理分析,这在第七章中讨论过,深度学习训练的生产化。该插件可以让你分析 PyTorch 推理代码,找出潜在的瓶颈。

使用 PTS 提供模型服务

让我们回顾一下如何使用 PTS 在 SageMaker 上部署 PyTorch 模型。我们将使用 HuggingFace 模型库中已在问答 NLP 任务上训练的 Distilbert 模型。示例代码可以在这里找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter9/2_PyTorch_Torchserve.ipynb

在 SageMaker 上为 PTS 打包模型

在 SageMaker 上使用 PTS 模型服务器时,你可以选择以下两种选项之一:

  • 使用 Python SageMaker SDK 中的 PyTorchModel 类部署你的模型。在这种情况下,你的模型归档只需要提供必要的模型工件(例如模型权重、查找表、分词器等)。作为 PyTorchModel 对象配置的一部分,你将提供推理代码和其他依赖项,SageMaker 会自动为 PTS 打包它。

  • 你也可以将模型与推理代码一起打包在一个归档文件中。虽然这种方法需要额外的工作,但它允许你创建模型包并部署模型,而无需使用 SageMaker SDK。在这种情况下,SageMaker 期望以下目录结构:

    model.tar.gz/
    |- model_weights.pth
    |- other_model_artifacts
    |- code/
      |- inference.py
      |- requirements.txt  # optional
    

在这个示例中,我们将使用第一个选项:

  1. 以下 Bash 脚本将下载所需的 HuggingFace 模型工件,并将它们打包成一个 tar.gz 归档 文件:

    mkdir distilbert-base-uncased-distilled-squad
    wget https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/pytorch_model.bin -P distilbert-base-uncased-distilled-squad
    wget https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/tokenizer.json -P distilbert-base-uncased-distilled-squad
    wget https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/tokenizer_config.json -P distilbert-base-uncased-distilled-squad
    wget https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/vocab.txt -P distilbert-base-uncased-distilled-squada
    wget https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json -P distilbert-base-uncased-distilled-squad
    tar -C "$PWD" -czf distilbert-base-uncased-distilled-squad.tar.gz  distilbert-base-uncased-distilled-squad/
    
  2. 然后,我们需要使用以下代码将模型归档上传到 Amazon S3:

    import sagemaker
    from sagemaker import get_execution_role
    sagemaker_session = sagemaker.Session()
    role = get_execution_role()
    bucket = sagemaker_session.default_bucket()
    prefix = 'torchserve'
    s3_path = 's3://{}/{}'.format(bucket, prefix)
    model_data = sagemaker_session.upload_data('distilbert-base-uncased-distilled-squad.tar.gz',bucket,os.path.join(prefix, 'model-artifacts'))
    

接下来,我们需要准备一些代码,从上传的模型工件中加载模型并执行推理和数据处理。这段代码在 PTS 术语中被称为 推理处理程序

准备推理处理程序

SageMaker 要求您提供一些代码以加载模型并运行预测,以便您可以预处理传入的推理请求并后处理响应。为了执行这些操作,您需要实现 model_fn()predict_fn()input_fn()output_fn() 方法。您可以在此找到使用 HuggingFace Pipeline API 的推理处理器实现:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter9/2_src/pipeline_predictor.py

将模型部署到 SageMaker 端点

使用 SageMaker SDK 在 PTS 上部署模型非常简单。为了配置 PTS,我们可以使用 "env" 字典在服务容器中设置适当的环境变量。请注意,在这里,我们通过 "entry_point" 参数显式引用推理代码。按照以下步骤操作:

  1. 作为前提条件,您可以将任何其他依赖项(例如自定义库或 requirements.txt)添加到 "source_dir" 位置。SageMaker SDK 将自动将这些资产与模型数据合并到 PTS 所需的 MAR 存档中:

    from sagemaker.pytorch import PyTorchModel
    env = {
        "SAGEMAKER_TS_BATCH_SIZE": "2",
        "SAGEMAKER_TS_MAX_BATCH_DELAY": "1000",
        "SAGEMAKER_TS_RESPONSE_TIMEOUT" : "120",
        "SAGEMAKER_TS_MIN_WORKERS" : "1",
        "SAGEMAKER_TS_MAX_WORKERS" : "2"
        }
    model = PyTorchModel(model_data=model_data,
                       role=role, 
                       entry_point='pipeline_predictor.py',
                       source_dir='2_src',
                       framework_version='1.9.0',
                       py_version='py38',
                       env=env,
                       sagemaker_session=sagemaker_session)
    
  2. 现在,我们可以定义请求/响应对的端点配置以及支持的序列化器和反序列化器:

    from sagemaker.serializers import JSONSerializer
    from sagemaker.deserializers import JSONDeserializer
    remote_predictor = model.deploy(initial_instance_count=1, instance_type="ml.g4dn.4xlarge", serializer=JSONSerializer(), deserializer=JSONDeserializer())
    
  3. 现在,我们可以通过调用 .predict() 方法进行预测:

    remote_predictor.predict(data)
    
  4. 我们还可以确认我们的 PTS 配置已正确应用。为此,您可以打开 SageMaker 端点日志流,并搜索如下所示的日志行:

    Model config: 
    { "model": { "1.0": { "defaultVersion": true, "marName": "model.mar", "minWorkers": 1, "maxWorkers": 2, "batchSize": 3, "maxBatchDelay": 100000, "responseTimeout": 120 } } }
    

在本节中,我们讨论了如何使用 PTS 来服务 PyTorch 模型。在实际生产系统中,您可能更愿意将模型转换为 TorchScript 格式,并进一步尝试批处理和工作节点扩展选项,以优化您的特定用例需求。

在下一节中,我们将回顾一个功能丰富的框架无关的模型服务器,称为 NVIDIA Triton。

使用 NVIDIA Triton

NVIDIA Triton 是由 NVIDIA 开发的一个开源模型服务器。它支持多种深度学习框架(如 TensorFlow、PyTorch、ONNX、Python 和 OpenVINO),以及各种硬件平台和运行时环境(NVIDIA GPUs、x86 和 ARM CPUs,以及 AWS Inferentia)。Triton 可用于云和数据中心环境中的推理,以及边缘设备或移动设备。Triton 在多种 CPU 和 GPU 平台上进行了性能和可扩展性优化。NVIDIA 提供了一种专门的工具,用于性能分析和模型分析,以提高 Triton 的性能。

与 SageMaker 的集成

你可以通过使用预构建的 SageMaker DL 容器来使用 Triton 模型服务器。请注意,SageMaker Triton 容器不是开源的。你可以在这里找到最新的 Triton 容器列表:github.com/aws/deep-learning-containers/blob/master/available_images.md#nvidia-triton-inference-containers-sm-support-only

SageMaker 在部署模型到 Triton 时不要求提供推理自定义代码。然而,你需要为每个你打算提供服务的模型提供一个 Triton config.pbtxt文件。该配置指定了推理请求/响应对的 API 契约以及模型如何提供服务的其他参数。你可以通过阅读官方 Triton 文档来查看可能的配置参数:github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md

同时需要注意的是,与 TFS 和 PTS 不同,写本文时,SageMaker 尚不支持在 Triton 上托管多个独立的模型。但是,你仍然可以有同一个模型的多个版本,或者将多个模型组织成一个管道。

优化 Triton 推理

Triton 提供多种工具来提高你的性能:

  • 模型分析器允许你了解模型的 GPU 内存使用情况,以便你了解如何在单个 GPU 上运行多个模型。

  • 性能分析器允许你分析 Triton 的推理和吞吐量。

你无法直接在 SageMaker Triton Endpoint 上运行性能分析器,因为 SageMaker 推理 API 与 Triton 推理 API 不匹配。为了解决这个限制,你可以在具有目标硬件加速器的 SageMaker Notebook 实例上本地运行 Triton 容器,并对其进行分析。

Triton 提供以下优化功能:

  • 动态批处理:将多个推理请求合并到一个批处理中,以提高 Triton 的吞吐量。这个功能类似于我们为 TFS 和 PTS 模型服务器讨论的批处理。

  • 模型实例:指定每个模型将有多少副本可用于推理。默认情况下,加载一个模型实例。拥有多个模型副本通常会提高延迟/吞吐量,因为它允许你将内存传输操作(例如,CPU 与 GPU 之间的传输)与推理计算重叠。拥有多个实例也能更有效地使用所有可用的 GPU 资源。

两个参数可以通过config.pbtxt文件进行配置。让我们在 SageMaker 上获得一些使用 Triton 的实际经验。

在 SageMaker 上使用 Triton 提供模型服务

在这个示例中,我们将使用 Triton 部署图像分类 PyTorch ResNet50 模型。我们的目标硬件加速器将是ml.g4dn实例。首先,我们需要将模型编译为 TensorRT 运行时;然后,编译后的模型将被打包并部署到 Triton 模型服务器。示例代码可以在此处找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter9/3_NVIDIA_Triton_Server.ipynb

请注意,以下小节中描述的模型编译过程是针对 PyTorch 框架的。如果你选择使用 TensorFlow 模型,则你的模型编译和配置会有所不同。你可以参考 Triton TensorFlow 后端的仓库获取详细信息:github.com/triton-inference-server/tensorflow_backend

为 Triton 编译模型

有几种方法可以将你的 PyTorch 模型编译为 TensorRT 格式,例如将 PyTorch 模型转换为 ONNX 格式。另一种方法是使用 PyTorch JIT 编译器将你的 eager 模型本地转换为 TorchScript 格式。最近,PyTorch 和 NVIDIA 团队实施了一种优化方式,使用Torch-TensorRT 编译器将 PyTorch 模型编译为 TensorRT 运行时。这个方法有几个优点,它允许你使用 TensorRT 特定的优化,例如 GP16 和 INT8 精度降低类型以及 NVIDIA GPU 权重稀疏性:

图 9.3 – 使用 TensorRT-Torch 编译 PyTorch 模型

图 9.3 – 使用 TensorRT-Torch 编译 PyTorch 模型

要使用 TensorRT-Torch 编译 PyTorch 模型,我们需要两个组件:

  • 用于编译的运行时环境。强烈建议使用 NVIDIA 的最新 PyTorch 容器来完成此操作。请注意,你需要在具有 NVIDIA GPU 的实例上运行此容器。例如,你可以在g4dn类型的 SageMaker Notebook 上运行此示例。

  • 编译代码。此代码将在 NVIDIA PyTorch Docker 容器内执行。

现在,让我们回顾一下编译代码:

  1. 我们将首先从 PyTorch Hub 加载模型,将其设置为评估模式,并将其放置在 GPU 设备上:

    import torch
    import torch_tensorrt
    import os
    torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
    MODEL_NAME = "resnet50"
    MODEL_VERSION = "1"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # load model
    model = (torch.hub.load("pytorch/vision:v0.10.0", MODEL_NAME, pretrained=True).eval().to(device))
    
  2. 接下来,我们将使用 TensorRT-Torch 编译器进行编译。在编译器配置中,我们将指定预期的输入和目标精度。请注意,由于我们计划使用动态批处理进行模型处理,因此我们将提供多个具有不同批次维度值的输入形状:

    # Compile with Torch TensorRT;
    trt_model = torch_tensorrt.compile(
        model,
        inputs=[
            torch_tensorrt.Input(
                min_shape=(1, 3, 224, 224),
                opt_shape=(8, 3, 224, 224),
                max_shape=(16, 3, 224, 224),
                dtype=torch.float32,
            )
        ],
        enabled_precisions={ torch.float32 },
    )
    
  3. 最后,我们将把模型保存到磁盘:

    # Save the model
    model_dir = os.path.join(os.getcwd(), "3_src", MODEL_NAME, MODEL_VERSION)
    os.makedirs(model_dir, exist_ok=True)
    print(model_dir)
    torch.jit.save(trt_model, os.path.join(model_dir, "model.pt"))
    
  4. 要执行此脚本,您需要通过以下命令启动一个 Docker 容器:docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm -v $PWD/chapter9/3_src:/workspace/3_src nvcr.io/nvidia/pytorch:22.05-py3

  5. 您的控制台会话将在容器内打开,在这里您可以通过运行python 3_src/compile_tensorrt.py命令执行编译脚本。

生成的model.pt文件将可在 Docker 容器外部的3_src目录中找到。

准备模型配置

之前,我们提到 Triton 使用一个特定约定的配置文件来定义模型签名和运行时配置。以下代码是一个config.pbtxt文件,我们可以用它来托管 ResNet50 模型。在这里,我们定义了批处理参数(最大批量大小和动态批处理配置)、输入和输出签名、模型副本以及目标硬件环境(通过instance_group对象):

name: "resnet50"
platform: "pytorch_libtorch"
max_batch_size : 128
input [
  {
    name: "input__0"
    data_type: TYPE_FP32
    dims: [ 3, 224, 224 ]
  }
]
output [
  {
    name: "output__0"
    data_type: TYPE_FP32
    dims: [ 1, 1000 ,1, 1]
  }
]
dynamic_batching {
   preferred_batch_size: 128
   max_queue_delay_microseconds: 1000
 }
instance_group {
  count: 1
  kind: KIND_GPU
}

详细了解 Triton 配置,请参考:github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md

打包模型工件

为了部署带有其配置的编译模型,我们需要将所有内容打包成一个tar.gz压缩包,并上传到 Amazon S3。以下代码展示了模型归档中的目录结构:

resnet50
|- 1
  |- model.pt
|- config.pbtxt

一旦模型包上传到 Amazon S3,我们可以部署我们的 Triton 端点。

部署 Triton 端点

Triton 推理容器不被 SageMaker Python SDK 支持。因此,我们需要使用 boto3 SageMaker 客户端来部署模型。请按照以下步骤操作:

  1. 首先,我们需要识别正确的 Triton 镜像。使用以下代码根据您使用的 Triton 服务器版本(我们使用了22.05用于模型编译和服务)以及您的 AWS 区域来查找 Triton 容器的 URI:

    account_id_map = {
     # <REDACTED_FOR_BREVITY>
    }
    region = boto3.Session().region_name
    if region not in account_id_map.keys():
        raise("UNSUPPORTED REGION")
    base = "amazonaws.com.cn" if region.startswith("cn-") else "amazonaws.com"
    triton_image_uri = "{account_id}.dkr.ecr.{region}.{base}/sagemaker-tritonserver:22.05-py3".format(
        account_id=account_id_map[region], region=region, base=base)
    
  2. 接下来,我们可以创建模型,定义模型数据和服务容器,以及其他参数,如环境变量:

    unique_id = time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime())
    sm_model_name = "triton-resnet50-" + unique_id
    container = {
        "Image": triton_image_uri,
        "ModelDataUrl": model_data,
        "Environment": {"SAGEMAKER_TRITON_DEFAULT_MODEL_NAME": "resnet50"},
    }
    create_model_response = sm_client.create_model(
        ModelName=sm_model_name, ExecutionRoleArn=role, PrimaryContainer=container
    )
    
  3. 然后,我们可以定义端点配置:

    endpoint_config_name = "triton-resnet50-" + unique_id
    create_endpoint_config_response = sm_client.create_endpoint_config(
        EndpointConfigName=endpoint_config_name,
        ProductionVariants=[
            {
                "InstanceType": "ml.g4dn.4xlarge",
                "InitialVariantWeight": 1,
                "InitialInstanceCount": 1,
                "ModelName": sm_model_name,
                "VariantName": "AllTraffic",
            }
        ],)
    
  4. 现在,我们准备好部署我们的端点:

    endpoint_name = "triton-resnet50-" + unique_id
    create_endpoint_response = sm_client.create_endpoint(
        EndpointName=endpoint_name, EndpointConfigName=endpoint_config_name)
    

一旦端点部署完毕,您可以检查 SageMaker 的端点日志,以确认 Triton 服务器已启动并且模型已成功加载。

运行推理

为了运行推理,我们必须根据config.pbtxt中定义的模型签名构造有效负载。请查看以下推理调用。响应也将遵循已定义的输出签名:

payload = {
    "inputs": [
        {
            "name": "input__0",
            "shape": [1, 3, 224, 224],
            "datatype": "FP32",
            "data": get_sample_image(),
        }
    ]
}
response = runtime_sm_client.invoke_endpoint(   EndpointName=endpoint_name, ContentType="application/octet-stream", Body=json.dumps(payload))
predictions = json.loads(response["Body"].read().decode("utf8"))

本节描述了 Triton 模型服务器的基本功能以及如何在 Amazon SageMaker 上使用它。建议您参考 Triton 文档,以了解高级功能和优化技术。请记住,根据您选择的模型格式和深度学习框架,您的模型配置将有所不同。您可以查看 AWS 针对 BERT 模型的 Triton 服务器详细基准测试,链接地址为 aws.amazon.com/blogs/machine-learning/achieve-hyperscale-performance-for-model-serving-using-nvidia-triton-inference-server-on-amazon-sagemaker/。这些基准测试为您提供了实验和调优 Triton 托管模型的良好起点。

总结

在本章中,我们讨论了如何在 Amazon SageMaker 上使用流行的模型服务器——TensorFlow Serving、PyTorch TorchServe 和 NVIDIA Triton。每个模型服务器都提供了丰富的功能,用于部署和调优模型推理。选择特定模型服务器的原因可能由深度学习框架、目标硬件和运行时环境以及其他偏好决定。NVIDIA Triton 支持多种模型格式、目标硬件平台和运行时。同时,TensorFlow Serving 和 TorchServe 提供与各自深度学习框架的原生集成。无论选择哪个模型服务器,为了确保计算资源和推理性能的最佳利用,建议规划如何通过各种服务器配置对模型进行负载测试和基准测试。

在下一章,第十章推理工作负载的操作化,我们将讨论如何在生产环境中迁移和管理推理工作负载。我们将回顾 SageMaker 在优化推理工作负载成本、进行 A/B 测试、根据推理流量模式进行端点资源的弹性伸缩,以及多模型和多容器端点等高级部署模式方面的能力。

第十章:操作化推理工作负载

第八章《考虑推理硬件》和第九章《实现模型服务器》中,我们讨论了如何在 Amazon SageMaker 上设计你的深度学习DL)推理工作负载。我们还回顾了如何为推理工作负载选择适当的硬件、优化模型性能并根据特定的使用案例要求调整模型服务器。在本章中,我们将重点讨论如何在推理工作负载部署到测试和生产环境后,进行操作化。

在本章中,我们将首先回顾一些高级模型托管选项,如多模型多容器无服务器推理端点,以优化你的资源利用率和工作负载成本。接着,我们将介绍 SageMaker 的应用自动扩展服务,它提供了另一种提高资源利用率的机制。自动扩展使你能够动态地将推理流量需求与预置的推理资源匹配。

接下来,我们将讨论如何在不影响终端用户的情况下,持续推广模型和模型版本。我们还将介绍一些用于 A/B 测试和模型候选质量保证的高级部署模式。为此,我们将回顾 SageMaker 的模型变体部署安全边界功能。

然后,我们将回顾如何使用 SageMaker 的模型监控器来监控模型和推理数据的质量。本章的最后,我们将讨论如何根据你的使用案例类型、业务需求和技术要求选择最佳的推理工作负载配置。

本章将涵盖以下内容:

  • 管理推理部署

  • 监控推理工作负载

  • 选择工作负载配置

本章结束时,你将掌握如何操作化 SageMaker 推理工作负载的理解和实践技能。

技术要求

在本章中,我们将提供代码示例,帮助你开发实践技能。完整的代码示例可以在这里查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter10/

要跟随本代码,你需要以下条件:

  • 一个 AWS 账户和具有管理 Amazon SageMaker 资源权限的 IAM 用户。

  • 一个 SageMaker 笔记本、SageMaker Studio 笔记本或已建立的本地 SageMaker 兼容环境。

  • 在你的 AWS 账户中访问 GPU 训练实例。本章中的每个示例将提供推荐使用的实例类型。你可能需要增加你的计算配额,以便启用 GPU 实例进行SageMaker 训练作业。在这种情况下,请按照docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.xhtml上的说明进行操作。

  • 你需要通过运行pip install -r requirements.txt来安装所需的 Python 库。包含所需库的文件可以在chapter10目录中找到。

  • 在本章中,我们将提供编译模型以进行推理的示例,这需要访问特定的加速器类型。请作为模型服务器示例的一部分,查看实例推荐。

管理推理部署

第一章与 Amazon SageMaker 一起介绍深度学习中,我们讨论了 SageMaker 在运行推理工作负载时提供的几种选项,这些选项取决于你的用例需求,具体如下:

  • 实时端点设计用于具有低延迟要求的推理用例。它在有效载荷大小(最大 5MB)和响应延迟(最多 60 秒)上有一定的限制。

  • 批量转换作业是处理大规模批量推理请求的离线选项。

  • 异步端点允许你排队并以接近实时的方式处理推理请求。与实时端点相比,它对推理有效载荷大小有更高的限制(最大 1GB)。

到目前为止,在本书中我们已经介绍了如何为推理工作负载部署单一模型。所有三种推理选项都支持这一方法。

然而,对于实时端点,有可能将多个模型和模型版本(称为生产 变体)打包并部署在一个单一端点后面。在本节中,我们将深入探讨这些模型部署策略,并重点介绍实施细节、其优势及某些限制。

此外,我们还将回顾最近推出的无服务器推理端点。与实时端点类似,无服务器端点旨在实时服务用户。然而,在无服务器端点的情况下,你可以使用计算资源,而无需选择并扩展推理实例。

考虑模型部署选项

在许多情况下,将单个模型托管在专用的 SageMaker 实时端点后面可能导致资源利用率低下并产生额外的费用,这些费用是可以避免的。例如,当你需要同时托管一批模型,每个模型的资源需求都很低时,将每个模型托管在单独的端点后面将是一项重大且可避免的成本。

SageMaker 提供了一系列模型部署选项,可以解决更复杂的用例。在以下小节中,我们将讨论它们的目标用例、优点和限制。

多模型端点

多模型端点MME)是一种特殊类型的 SageMaker 模型端点,允许您在单个端点后同时托管数千个模型。这种类型的端点适用于模型大小相似、资源要求较低、可以从同一推理容器中提供服务的场景。

MME 及其底层模型服务器管理资源分配,例如在实例内存不足时卸载不常用的模型并加载请求的模型。这会导致当用户请求当前未加载到内存中的模型时,出现额外的推理延迟。因此,MME 可能不适用于对低延迟要求持续较低的场景。当托管大型模型且流量模式均匀分布时,这种额外的延迟可能会增加,因为这将导致频繁地卸载和加载模型。

要配置 MME,您需要将每个模型(模型工件和推理代码)打包成单独的归档文件,并将其上传到 Amazon S3。MME 实例配置完毕后,它将从 S3 位置下载到实例磁盘,将模型加载到实例内存中。默认情况下,如果 MME 的实例磁盘空间和/或实例内存不足,SageMaker 会从本地磁盘中删除最不常用的模型和/或从内存中卸载模型,以便为请求的模型腾出空间。

下图展示了 MME 架构:

图 10.1 – MME 架构

图 10.1 – MME 架构

MME 支持 PyTorch 和 TensorFlow 推理容器。您还可以根据推理流量自动扩展和缩减 MME。MME 允许您直接调用模型以及由多个模型组成的推理管道。

在选择实例类型和系列时,请考虑以下几个方面:

  • 实例内存定义了可以同时加载多少个模型。

  • 实例磁盘大小定义了可以在本地缓存多少个模型,从而避免从 S3 进行昂贵的下载操作。

  • vCPU 的数量定义了可以同时处理多少个推理请求。

请注意,基于 GPU 的实例不支持 MME,这限制了可以在合理的服务水平协议(SLA)内使用 MME 托管的模型架构。

现在,让我们学习如何实现 MME。

实现 MME

在这个代码示例中,我们将学习如何使用 MME 同时部署两个 NLP 模型。一个模型分析德语文本的情感,另一个模型分析英语文本的情感。我们将使用 HuggingFace PyTorch 容器来完成此任务。完整的代码可以在此处查看:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter10/1_Multi_Model_Endpoint.ipynb

对于此任务,我们将使用两个模型,分别训练用于预测英语和德语文本的情感:distilbert-base-uncased-finetuned-sst-2-englisholiverguhr/german-sentiment-bert。请按照以下步骤进行:

  1. 我们将首先从 HuggingFace 模型库中获取模型并将其保存在本地。以下代码展示了英语模型:

    import torch
    from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
    en_tokenizer = DistilBertTokenizer.from_pretrained(EN_MODEL)
    en_model = DistilBertForSequenceClassification.from_pretrained(EN_MODEL)
    en_model_path = "models/english_sentiment"
    os.makedirs(en_model_path, exist_ok=True)
    en_model.save_pretrained(save_directory=en_model_path)
    en_tokenizer.save_pretrained(save_directory=en_model_path)
    

结果,将下载以下工件:

('models/english_sentiment/tokenizer_config.json',
 'models/english_sentiment/special_tokens_map.json',
 'models/english_sentiment/vocab.txt',
 'models/english_sentiment/added_tokens.json')

这些模型工件稍后将被添加到模型数据包中。但首先,我们需要开发推理脚本。

  1. MME 与单模型端点的推理脚本具有相同的要求。以下代码展示了英语模型的推理脚本,脚本实现了模型加载、推理以及数据的前后处理所需的方法:

    def model_fn(model_dir):
        tokenizer = DistilBertTokenizer.from_pretrained(model_dir)
        model = DistilBertForSequenceClassification.from_pretrained(model_dir)
        return model, tokenizer
    def input_fn(serialized_input_data, content_type=JSON_CONTENT_TYPE):
        if content_type == JSON_CONTENT_TYPE:
            input_data = json.loads(serialized_input_data)
            return input_data
        else:
            Exception("Requested unsupported ContentType in Accept: " + content_type)
    def predict_fn(input_data, model_tokenizer_tuple):
        model, tokenizer = model_tokenizer_tuple
        inputs = tokenizer(input_data, return_tensors="pt")
        with torch.no_grad():
            logits = model(**inputs).logits
        predicted_class_id = logits.argmax().item()
        predictions = model.config.id2label[predicted_class_id]
        return predictions
    def output_fn(prediction_output, accept=JSON_CONTENT_TYPE):
        if accept == JSON_CONTENT_TYPE:
            return json.dumps(prediction_output), accept
        raise Exception("Requested unsupported ContentType in Accept: " + accept)
    
  2. 接下来,我们需要为 MME 打包模型和推理代码。SageMaker 要求一个特定的目录结构,该结构在 PyTorch 和 TensorFlow 容器中有所不同。对于 PyTorch 容器,模型和代码应打包成一个单独的tar.gz压缩包,并具有以下结构:

    model.tar.gz/
                 |- model.pth # and any other model artifacts
                 |- code/
                         |- inference.py
                         |- requirements.txt # optional
    

每个模型应有一个模型包。一旦包在本地准备好,我们需要将它们上传到 Amazon S3,并保存相应的 URI:

en_model_data = sagemaker_session.upload_data('models/english_sentiment.tar.gz', bucket=bucket,key_prefix=prefix)
ger_model_data = sagemaker_session.upload_data('models/german_sentiment.tar.gz', bucket=bucket,key_prefix=prefix)
  1. 一旦数据上传完成,我们需要定义相应的服务容器,并配置它以供 MME 使用。以下代码根据所需的运行时配置和任务(推理)定位 PyTorch 容器:

    from sagemaker import image_uris
    HF_VERSION = '4.17.0'
    PT_VERSION = 'pytorch1.10.2'
    pt_container_uri = image_uris.retrieve(framework='huggingface',
                                    region=region,
                                    version=HF_VERSION,
                                    image_scope='inference',
                                    base_framework_version=PT_VERSION,
                                    instance_type='ml.c5.xlarge')
    
  2. 然后,我们需要配置 MME 参数。具体来说,我们必须定义MultiModel模式。请注意,我们提供了两个特定的环境变量 – SAGEMAKER_PROGRAMSAGEMAKER_SUBMIT_DIRECTORY – 以便 SageMaker 推理框架知道如何注册模型处理器:

    container  = {
        'Image': pt_container_uri,
        'ContainerHostname': 'MultiModel',
        'Mode': 'MultiModel',
        'ModelDataUrl': mm_data_path,
        'Environment': {
        'SAGEMAKER_PROGRAM':'inference.py',
        'SAGEMAKER_SUBMIT_DIRECTORY':mm_data_path
        }
    }
    
  3. 配置 MME 的最后一步是创建 SageMaker 模型实例、端点配置和端点本身。在创建模型时,我们必须提供前一步中启用了MultiModel的容器。为了简便起见,我们省略了端点配置和端点的创建:

    unique_id = datetime.datetime.now().strftime("%Y-%m-%d%H-%M-%S")
    model_name = f"mme-sentiment-model-{unique_id}"
    create_model_response = sm_client.create_model(
        ModelName=model_name,
        PrimaryContainer=container,
        ExecutionRoleArn=role,
    )
    
  4. 一旦端点创建完成,我们就可以运行并调用我们的模型。为此,在调用请求中,我们需要提供一个名为TargetModel的特殊参数,如下所示:

    ger_response = runtime_sm_client.invoke_endpoint(
        EndpointName=endpoint_name,
        ContentType="application/json",
        Accept="application/json",
        TargetModel="german_sentiment.tar.gz",
        Body=json.dumps(ger_input),
    )
    

虽然 MME 功能提供了一种便捷的方式来优化运行多个相似模型时的推理成本,但它要求模型具有相同的运行时环境(换句话说,它们必须使用相同的推理容器)。为了解决需要在不同推理容器中托管多个模型的场景,SageMaker 支持多容器端点MCEs),如下一节所示。

多容器端点

MCE 允许你同时托管最多 15 个推理容器。在这种情况下,每个容器将服务于自己的模型。MCE 非常适合用于需要不同运行时环境/容器的模型场景,但并非每个模型都能充分利用可用的实例资源。另一种场景是当模型在不同时间被调用时。

与 MME 不同,MCE 不会根据容器的调用模式缓存或卸载容器。因此,你需要确保推理容器总共拥有足够的资源来在端点实例上运行。如果实例资源(例如,实例内存)不足以运行所有容器,可能会在创建 MCE 时出现错误。因此,在选择实例配置时,你需要考虑所有推理容器的总资源需求。每个推理容器将为其提供按比例分配的资源。以下图显示了 MCE 架构:

图 10.2 – MCE 架构

图 10.2 – MCE 架构

你可以自动扩展 MCE。它支持Direct模式(直接调用推理容器)或Serial模式(按顺序调用多个容器)。

在写这本书时,MCEs 不支持基于 GPU 的实例。

现在,让我们通过一个简单的例子来学习如何创建 MCE,运行 TensorFlow 和 PyTorch 模型并行工作。这将让你获得一些实际技能,了解如何创建和使用 MCE。

实现 MCE

在这个例子中,我们将使用不同运行时环境的两个 NLP 模型运行推理工作负载:TensorFlow 和 PyTorch。我们将在 TensorFlow 容器中托管问答模型,在 PyTorch 容器中托管文本摘要模型。

创建 MCE 与创建 MME 非常相似,除了几个显著的不同之处,我们将在以下步骤中进行突出说明:

  1. 获取模型数据、推理脚本和模型包装的过程与我们为 MME 所做的一致。请注意,由于我们的一个端点将运行 TensorFlow 容器,因此问答模型应遵循以下目录结构:

    model.tar.gz/
                 |--[model_version_number]/
                                           |--variables
                                           |--saved_model.pb
                code/
                    |--inference.py
                    |--requirements.txt # optional
    
  2. 接下来,我们将配置容器并创建模型包。请注意,在创建模型包时,我们提供了两个容器和端点模式Direct

    model_name = f"mce-nlp-model-{unique_id}"
    create_model_response = sm_client.create_model(
        ModelName=model_name,
        Containers=[tensorflow_container, pytorch_container],
        InferenceExecutionConfig={"Mode": "Direct"},
        ExecutionRoleArn=role,
    )
    
  3. 然后,我们将创建端点配置和端点。这个步骤类似于 MME,因此我们省略了代码片段以简化内容。

  4. 一旦端点部署完成,我们就可以开始发送推理流量。请注意,我们提供了 TargetContainerHostname 头信息,以便 SageMaker 知道将推理请求路由到哪里:

    tf_response = runtime_sm_client.invoke_endpoint(
        EndpointName=endpoint_name,
        ContentType="application/json",
        Accept="application/json",
        TargetContainerHostname="tensorflow-distilbert-qa",
        Body=json.dumps(qa_inputs),
    )
    

到目前为止,我们已经讨论了如何在 SageMaker 上托管多个模型。接下来,我们将讨论如何在保持端点对最终用户可用的同时,安全地推广新版本的模型(或完全不同的模型)。为此,我们将回顾 SageMaker 多变体端点。

多变体端点

生产变体是 SageMaker 特有的概念,它定义了模型、其容器以及运行该模型所需资源的组合。因此,这是一个极其灵活的概念,可以用于不同的使用场景,例如以下情况:

  • 拥有相同运行时和资源需求的不同模型版本

  • 拥有不同运行时和/或资源需求的不同模型

  • 相同模型与不同的运行时和/或资源需求

此外,作为变体配置的一部分,你还需要定义其流量权重,这些权重可以在不影响端点可用性的情况下进行更新。部署后,生产变体可以直接调用(从而绕过 SageMaker 流量控制),也可以作为 SageMaker 端点调用的一部分(此时不会绕过 SageMaker 流量控制)。以下图示提供了更多细节:

图 10.3 – 使用流量控制的生产变体(左)和直接调用的生产变体(右)

图 10.3 – 使用流量控制的生产变体(左)和直接调用的生产变体(右)

在更新生产变体时,实时端点保持可用,并且不会对最终用户造成中断。这也意味着你将产生额外的费用,因为每个生产变体都会有一个关联的费用。

现在,让我们看看如何使用生产变体来测试一个新的生产变体。

使用生产变体进行 A/B 测试

在这个示例中,我们将为相同的问答 NLP 任务注册两个不同的模型。然后,我们将使用生产变体权重对推理流量进行流量控制,并直接调用模型。完整代码可在此处获取:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter10/4_AB_Testing.ipynb。请按照以下步骤进行操作:

  1. 我们将从使用 HuggingFaceModel 类创建两个 HuggingFace 模型开始。为了简洁起见,我们省略了这部分。

  2. 然后,我们将创建两个不同的端点变体。我们从相等权重参数开始,这告诉 SageMaker 推理流量应该在模型变体之间均匀分配:

    from sagemaker.session import production_variant
    variant1 = production_variant(
        model_name=model1_name,
        instance_type="ml.c5.4xlarge",
        initial_instance_count=1,
        variant_name="Variant1",
        initial_weight=1,
    )
    variant2 = production_variant(
        model_name=model2_name,
        instance_type="ml.c5.4xlarge",
        initial_instance_count=1,
        variant_name="Variant2",
        initial_weight=1,
    )
    
  3. 之后,我们根据配置的生产变体创建端点:

    from datetime import datetime
    endpoint_name = f"ab-testing-{datetime.now():%Y-%m-%d-%H-%M-%S}"
    sagemaker_session.endpoint_from_production_variants(
        name=endpoint_name, production_variants=[variant1, variant2]))
    
  4. 一旦端点部署完成,我们就可以对新创建的端点进行推理。运行以下代码后,结果统计应显示每个生产变体大约服务了 ~50% 的推理流量:

    results = {"Variant1": 0, "Variant2": 0, "total_count": 0}
    for i in range(20):
        response = sm_runtime_client.invoke_endpoint(EndpointName=endpoint_name, ContentType="application/json", Body=json.dumps(data))
        results[response['InvokedProductionVariant']] += 1
        results["total_count"] += 1
    
  5. 接下来,我们可以更新端点变体的权重。重新运行先前的推理测试循环,现在应该显示只有 ~10% 的流量由 "Variant1" 服务,这与提供的变体流量权重一致:

    sm_client.update_endpoint_weights_and_capacities(
        EndpointName=endpoint_name,
        DesiredWeightsAndCapacities=[
            {"DesiredWeight": 10, "VariantName": "Variant1"},
            {"DesiredWeight": 90, "VariantName": "Variant2"},])
    
  6. 我们还可以绕过 SageMaker 的流量整形,直接通过使用 TargetVariant 参数调用特定变体,具体如下:

    sm_runtime_client.invoke_endpoint(EndpointName=endpoint_name, TargetVariant="Variant2", ContentType="application/json", Body=json.dumps(data))
    

SageMaker 的生产变体为您提供了一种灵活的机制,用于在生产或类生产环境中运行推理工作负载。

无服务器推理端点

使用无服务器推理端点SIEs)是 SageMaker 提供的另一种部署选项。它允许您在无需预配置和配置底层端点实例的情况下,提供实时推理端点。SageMaker 会根据推理流量自动配置和扩展底层可用计算资源。如果没有推理流量,您的 SIE 可以将其缩减为 0。

SIEs 非常适合流量模式不均匀的场景,且您可以容忍在冷启动期间出现短暂的延迟增加。冷启动时长指定了配置新的无服务器资源并部署模型运行时环境所需的时间。由于较大的模型通常比小模型需要更长的部署时间,因此它们的冷启动时长也会更长。无服务器推理的一个潜在用例是在测试和沙盒环境中使用它。使用 SIE,您只需为 SIE 处理推理请求所花费的时间付费。

无服务器推理在功能上类似于 SageMaker 实时推理。它支持多种推理容器类型,包括 PyTorch 和 TensorFlow 推理容器。然而, 无服务器推理也有一些限制,主要包括以下几点:

  • 没有可用的 GPU 资源

  • 实例磁盘大小为 5 GB

  • 端点的最大并发量为 200;超出此限制的请求将被 SageMaker 限流。

  • 冷启动时长取决于您的模型大小和推理容器的启动时间

在创建 SIE 资源时,您可以从可用的内存选项列表中进行选择,SageMaker 会自动分配相应数量的 vCPU。在配置内存时,您需要确保内存大小至少略大于模型大小,并且最小内存大小必须为 1,024 MB;最大内存为 6,144 MB。如果您的模型性能受 CPU 限制,您可以选择更大的内存配置,以获得更多的 vCPU 资源。

现在,让我们看看如何使用 SageMaker Python SDK 部署一个无服务器端点。

部署无服务器端点

在这个示例中,我们将部署来自 HuggingFace 模型库的问答 NLP 模型。完整代码可见这里:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter10/3_Serverless_Inference.ipynb。按照以下步骤操作:

  1. 我们将从定义要部署的 SageMaker 模型开始。为此,我们将使用内置的 image_uris 方法获取 HuggingFace 推理容器的 CPU 版本,如下所示:

    From sagemaker import image_uris
    HF_VERSION = '4.17.0'
    PT_VERSION = 'pytorch1.10.2'
    hf_container_uri = image_uris.retrieve(framework='huggingface',
                              region=region,
                                    version=HF_VERSION,
                                    image_scope='inference',
                                    base_framework_version=PT_VERSION,
                                    instance_type='ml.c5.xlarge')
    
  2. 接下来,我们将使用 HuggingFaceModel 实例来配置模型架构和目标 NLP 任务:

    Hub = {
        'HF_MODEL_ID':'distilbert-base-uncased-distilled-squad',
        'HF_TASK':'question-answering'
    }
    huggingface_model = HuggingFaceModel(
       env=hub,  
       role= role, 
       transformers_version=HF_VERSION,
       pytorch_version=PT_VERSION,     
       image_uri=hf_container_uri,     
    )
    
  3. 接下来,我们将定义无服务器配置并部署我们的第一个端点。此处,memory_size_in_mb 参数定义了端点背后的初始内存,max_concurrency 参数定义了端点在推理流量被 SageMaker 限流之前可以处理的最大并发调用数:

    Serverless_config = ServerlessInferenceConfig(
        memory_size_in_mb=4096, max_concurrency=10,
    )
    predictor = huggingface_model.deploy(
        serverless_inference_config=serverless_config
    )
    

就这样!几分钟后,你的端点将被部署。之后,你可以像使用任何其他实时端点一样使用它。

使用无服务器推理时,SageMaker 会在无需太多输入的情况下自动扩展端点,除了内存大小和并发性设置。在下一节中,我们将回顾端点自动扩展功能,它为你提供了更多细粒度的扩展行为控制。

高级模型部署技巧

在本节中,我们将讨论一些管理 SageMaker 推理资源的高级技巧,即自动扩展和蓝绿部署。

自动扩展端点

SageMaker 允许你自动扩展(增加实例数量)和收缩(减少实例数量)实时端点和异步端点。当推理流量增加时,扩展可以保持端点性能稳定,同时将成本控制在最低。流量减少时,收缩可以帮助你降低推理成本。对于实时端点,最小实例大小为 1;而异步端点则可以扩展到 0 实例。下图展示了这一点:

图 10.4 – 自动扩展概念

图 10.4 – 自动扩展概念

在扩展事件期间,SageMaker 端点始终保持完全可用,对终端用户可用。在收缩端点的情况下,SageMaker 会自动将流量从实例中排空,以便可以移除这些实例。为了确保更高的可用性,SageMaker 会将实例分布在不同的 可用区

要自动扩展你的端点,你需要为模型创建一个生产变体。之后,你必须在 自动扩展策略 中定义期望的扩展行为。SageMaker 支持四种类型的扩展策略,如下所示:

  • TargetTrackingScaling)允许你根据特定 Amazon CloudWatch 指标的值来扩展端点。SageMaker 默认支持多种端点指标,但你也可以使用自己的自定义指标。CPUUtilizationGPUUtilizationSageMakerVariantInvocationsPerInstance指标通常是一个不错的起点。

  • 阶梯扩展是一种更高级的扩展策略,允许你根据指标值变化的大小更精细地控制实例的配置数量。此策略需要根据不同的负载配置文件值进行仔细的配置和测试。

  • 定时扩展允许你根据预定的时间表扩展端点。例如,你可以在非工作时间缩减扩展,在工作高峰时段进行扩展。

  • 按需扩展根据用户的明确请求来改变端点实例的数量。

在选择和配置自动扩展策略时,你可以通过分析流量模式以及它们与端点指标的相关性来开始。负载配置文件定义了选择哪种类型的扩展策略,而与端点指标的相关性使你能够选择合适的跟踪指标。建议你从一个简单的基准开始(例如,使用 CPUUtilization 跟踪指标的简单扩展)。然后,你可以随着观察到其他流量模式以及自动扩展对它们的反应,逐步进行调整。

在以下示例中,我们将学习如何将自动扩展策略应用于 SageMaker 实时端点。

为推理端点实现自动扩展

在这个例子中,我们将学习如何将目标跟踪自动扩展策略应用于实时端点。完整的代码可以在这里找到:github.com/PacktPublishing/Accelerate-Deep-Learning-Workloads-with-Amazon-SageMaker/blob/main/chapter10/5_AutoScaling.ipynb。按照以下步骤操作:

  1. 我们将首先创建一个常规的 SageMaker 实时端点。为简便起见,我们省略了此部分代码。

  2. 接下来,我们将创建两个自动扩展资源:可扩展目标扩展策略。可扩展目标定义了我们希望通过应用程序自动扩展服务进行扩展的特定 AWS 资源。

在以下代码片段中,我们实例化了应用程序自动扩展服务的客户端,并将我们的 SageMaker 端点注册为可扩展目标。请注意,ResourceId参数定义了对特定端点和生产变体的引用。SageMaker 资源的ScalableDimension参数始终引用生产变体后面的实例数量。MinCapacityMaxCapacity定义了实例扩展的范围:

import boto3 
as_client = boto3.client('application-autoscaling') 
resource_id=f"endpoint/{predictor.endpoint_name}/variant/AllTraffic"
policy_name = f'Request-ScalingPolicy-{predictor.endpoint_name}'
scalable_dimension = 'sagemaker:variant:DesiredInstanceCount'
# scaling configuration
response = as_client.register_scalable_target(
    ServiceNamespace='sagemaker', #
    ResourceId=resource_id,
    ScalableDimension='sagemaker:variant:DesiredInstance Count', 
    MinCapacity=1,
    MaxCapacity=4
)
  1. 接下来,我们将为可扩展目标创建一个策略。在这里,我们选择使用目标跟踪策略类型,并设置以下参数:

    response = as_client.put_scaling_policy(
        PolicyName=policy_name,
        ServiceNamespace='sagemaker',
        ResourceId=resource_id,
        ScalableDimension=scalable_dimension,
        PolicyType='TargetTrackingScaling',
        TargetTrackingScalingPolicyConfiguration={
            'TargetValue': 10.0, # Threshold
            'PredefinedMetricSpecification': {
                'PredefinedMetricType': 'SageMakerVariantInvocationsPerInstance',
            },
            'ScaleInCooldown': 300, # duration until scale in
            'ScaleOutCooldown': 60 # duration between scale out
        }
    )
    
  2. 一旦策略就绪,我们可以进行测试。为此,我们需要生成足够的推理流量,超出目标指标值,并持续时间超过扩展冷却期。为此,我们可以使用 Locust.io 负载测试框架(locust.io/),它提供了一种简单的机制来模拟各种负载模式。按照笔记本中的说明创建 Locust 配置,为您的端点提供 AWS 凭证进行授权。

  3. 配置完成后,您可以启动 Locust 客户端,通过以下终端命令生成负载。它会在 5 分钟内生成最多 20 个并发用户的推理负载。这个负载配置应该会触发端点的扩展事件:

    locust -f ../utils/load_testing/locustfile.py --headless -u 20 -r 1 --run-time 5m
    
  4. 在负载测试期间,您可以在 Amazon CloudWatch 控制台中观察端点状态以及相关的扩展警报。首先,您可以看到已经根据提供的冷却期和目标指标值配置了扩展和收缩警报:

图 10.5 – SageMaker 端点的自动扩展警报

图 10.5 – SageMaker 端点的自动扩展警报

  1. 在初始扩展冷却期结束后,扩展警报会切换到告警中状态,从而导致端点进行扩展。请注意,在以下截图中,红线是跟踪指标的目标值,而蓝线是每个端点实例的调用次数:

图 10.6 – 触发了扩展警报

图 10.6 – 触发了扩展警报

  1. 在触发扩展后,您的端点状态将从in Service变为Updating。现在,我们可以运行describe_endpoint()方法来确认实例数量已经增加。由于我们在短时间内生成了足够大的并发负载,SageMaker 立即将端点扩展到最大实例数。以下代码用于describe_endpoint()方法:

      ...
      "ProductionVariants": [
        {
          "VariantName": "AllTraffic",
          ...
          ],
          "CurrentWeight": 1,
          "DesiredWeight": 1,
          "CurrentInstanceCount": 1,
          "DesiredInstanceCount": 4
        }
      ]
      'EndpointStatus': 'Updating'
      ...
    

由于我们不再运行推理流量生成器,一旦扩展收缩冷却期结束,我们应该预期端点会收缩。

在下一节中,我们将回顾如何使用 SageMaker 部署护栏安全可靠地部署模型候选。

使用蓝绿部署模式

到目前为止,我们已经讨论了如何通过 API 或 SDK 调用部署和更新 SageMaker 端点。然而,这种方法可能不适用于更新生产中的关键任务工作负载,因为在这种情况下,你需要额外的检查来确保生产部署的顺利进行。

SageMaker 部署保护措施是一个完全托管的端点推广机制。保护措施遵循蓝绿部署概念,这是 DevOps 实践中常见的一种方式。在这里,蓝色舰队是旧的部署(在 SageMaker 端点中是生产变体),而绿色舰队是要部署的新版本。SageMaker 会在蓝色舰队旁边配置绿色舰队。一旦绿色舰队准备就绪且健康,SageMaker 就会按照预定义的规则将流量从蓝色舰队转移到绿色舰队。

部署保护措施支持多种流量转移模式:

  • All at once模式一旦绿色舰队启动并健康,就一次性将所有流量从蓝色转移到绿色。在这一点上,SageMaker 会停用蓝色舰队。

  • Canary模式将一小部分流量转移到绿色舰队。如果金丝雀(canaries)健康,SageMaker 将剩余的流量转移到绿色舰队。之后,SageMaker 会停用蓝色舰队。

  • Linear模式逐步将流量从蓝色舰队转移到绿色舰队。

请注意,在蓝绿部署期间,您将为蓝色和绿色舰队的运行同时产生费用。如果在部署过程中绿色舰队变得不健康,SageMaker 会执行自动回滚到初始部署,以避免对最终用户体验产生影响。

部署保护措施不支持以下功能:

  • 市场容器

  • 多容器端点

  • 多模型端点

  • 多变体端点

  • 使用 Inferentia 基础实例的端点

  • 使用 Amazon SageMaker 模型监控器(启用数据捕获)的端点

部署保护措施的设置实践超出了本书的范围,因为这些任务通常由专门的 DevOps/MLOps 团队执行。然而,重要的是要了解 SageMaker 原生支持这些功能。

监控推理工作负载

在本节中,我们将介绍用于监控推理工作负载的可用机制。

使用 Amazon CloudWatch

本书中,我们多次提到 Amazon CloudWatch。SageMaker 依赖它来满足所有的监控需求,具体包括以下内容:

  • 使用 CloudWatch 日志来收集、组织和管理 SageMaker 日志(例如,您的模型服务器日志)。

  • 使用 CloudWatch 指标来衡量端点特征,例如延迟、资源利用率等。

  • 使用 CloudWatch 警报触发自动扩展事件。

SageMaker 推断工作负载支持多种出厂指标。根据选择的推断工作负载选项和部署模式,您的默认 SageMaker 指标可能会有所不同。例如,对于 MME,您将获得额外的默认指标以衡量一些特定特征,如模型性能和加载时间。我们建议您参考 SageMaker 文档,获取关于默认 SageMaker 指标的最新信息:docs.aws.amazon.com/sagemaker/latest/dg/monitoring-cloudwatch.xhtml

如果由于某种原因,现有的出厂指标不足以满足您的用例需求,您始终可以创建自定义指标。一些适用自定义指标的场景如下:

  • 您的模型和模型服务器需要自定义指标以进行适当的扩展。

  • 您需要更高分辨率的指标。请注意,SageMaker 默认指标为 1 秒分辨率。

  • 您需要进行自定义指标预处理。例如,您可能需要应用 CloudWatch 不支持的滑动窗口平均值。

您还可以创建自定义 CloudWatch 警报。请注意,您可以为指标和日志创建警报。CloudWatch 警报可用于通过电子邮件或短信通知您特定事件(这将需要将您的警报与 Amazon SNS 服务集成)。

另一个 CloudWatch 警报的热门用例是在触发警报后执行操作。我们已经看到 CloudWatch 警报如何用于调整您的 SageMaker 端点的大小。但是,您可以将警报用于任何其他自定义逻辑。例如,您可以将自定义警报与 Amazon Lambda 无服务器函数集成。一旦触发警报,将执行您的功能及其自定义逻辑(例如,端点更新操作)。

监控推断工作质量

SageMaker 模型监控是专为衡量和持续监控推断质量而设计的能力。它允许您计算推断输入和模型输出的基线统计数据,然后实时监控模型在与基线统计数据的比较中的表现。在显著偏离预定义统计约束条件的情况下,SageMaker 模型监控将生成警报,通知您模型可能未按期望的质量指标执行。

Model Monitor 包括多个组件,用于监控推断质量的不同方面:

  • 数据质量监控 允许您检测用于训练模型的数据与部署模型的实际推断流量之间的数据漂移。数据漂移通常会导致模型预测质量低于预期。为了检测数据漂移,Model Monitor 计算训练数据(基线)的统计数据,捕获推断流量,并持续比较这些统计数据与基线。

  • 模型质量监控使你能够将模型预测与预定义的真实标签进行比较。如果模型预测违反了预定义的真实标签约束,Model Monitor 会生成警报。

  • 偏差漂移监控使你能够检测模型预测中的偏差以及它随时间的变化。当推理流量与用于模型训练的数据不同,可能会引入模型偏差。为了检测偏差,Model Monitor 计算一个特定的偏差指标,称为预测标签中正比例差异DPPL)。当 DPPL 指标违反预定义的数值范围时,将生成警报。

  • 模型特征归因监控是确保在模型部署过程中不会引入新偏差的另一种方式。特征归因漂移意味着某一特定特征对推理结果的影响随着时间变化。

注意

Model Monitor 仅支持表格数据作为推理输入。这限制了其在深度学习推理中的适用性,因为在大多数情况下,深度学习模型用于对非结构化数据(例如图像或文本)进行推理。

有几种场景可以将 Model Monitor 应用于深度学习推理:

  • 如果你使用深度学习模型进行分类或回归推理任务,实际上这种情况很少发生,因为经典的机器学习ML)算法(例如 XGBoost)通常在此类任务上超越深度学习模型,并且相较于更昂贵的深度学习模型,它们在训练和推理时需要的资源也较少。

  • 如果你的推理输入可以在发送到 SageMaker 推理资源之前从非结构化格式转换为结构化格式——例如,如果你将非结构化文本转换为分词输入并发送进行推理。在这种情况下,分词后的输入可以表示为一个表格数据集,从而可以与 Model Monitor 一起使用。

请注意,你仍然可以使用 Model Monitor 来确保你的深度学习工作负载在分类或回归输出场景中的模型精度。

选择工作负载配置

在前面三章中,我们回顾了 Amazon SageMaker 提供的不同功能,用于设计和操作推理工作负载:从选择最佳计算实例和运行时环境到配置模型服务器以及管理和监控已部署的模型。

在本节中,我们将总结在选择推理工作负载配置时可以使用的各种选择标准。接着,我们将建议一个简单的算法来指导你在选择推理配置时的决策过程。

在设计推理工作负载时,你可以考虑以下选择标准:

  • 业务应用场景:这让你能够通过使用推理服务来了解你的商业机会和最终用户体验。分析你的应用场景能够推动一些重要的决策,比如选择合适的 SageMaker 推理选项和最终用户 SLA(服务水平协议)。

  • 推理 SLA:本书中我们讨论了两个关键的推理 SLA:延迟和吞吐量。了解期望的 SLA 可以推动一系列决策,比如选择使用哪种实例类型、模型服务器配置等。

  • 预算和成本:预测推理预算和设置监控预算使用的机制(推理运行成本)是非常重要的。如果预算超支,你可能需要有机制来应对此类事件(例如,发送通知、缩减端点规模等)。

  • 计算实例:选择计算实例时,你需要考虑多个因素,如你打算使用的模型架构、你的 SLA 等。选择实例类型的过程称为合理化(rightsizing),需要进行负载测试和基准测试。

  • 输入数据和推理流量:你需要了解你的数据大小(针对离线推理)和推理流量模式(针对在线推理)。例如,如果你的流量具有季节性模式,你可能能够利用端点自动扩展来最小化推理成本。

  • 模型运行时和部署:根据你的模型特性、推理流量模式以及选择的计算实例,你需要选择特定的 SageMaker 容器和模型打包配置(单一模型与多个模型通过一个端点进行部署)。另一个需要探索的方面是模型推广策略和生产环境中的质量保证。例如,在本章早些时候,我们讨论了如何使用生产变体组织在 SageMaker 实时端点上的 A/B 测试。

以下表格突出显示了可用的 SageMaker 推理选项的关键特性:

实时推理 批量转换 异步推理 无服务器推理
推理类型 在线(实时响应) 离线 在线(接近实时推理,冷启动时从 0 扩展) 在线(冷启动时从 0 扩展)
资源扩展 单个端点下的 1 到数百个实例 单个推理任务中的 1 到数百个实例 0 到数百个实例 0 到 200 个并发推理请求
负载大小 最大 6 MB 最大 100 MB 最大 1 GB 最大 4 MB
推理超时 60 秒 无限制 最长 15 分钟 60 秒
多模型/多容器支持
目标使用案例 当您需要具有一致的实时推理延迟时。支持广泛的计算实例和模型服务器。 当输入数据集可预先提供时,进行离线推理或处理。 当您需要处理较大的负载大小和/或处理时间,并且可以接受额外的推理延迟时。在没有推理流量时,可以进行零扩展来节省成本。 当您需要最低的管理开销和相关成本的实时推理时。仅为服务的推理请求付费,且可以扩展至零。

图 10.7 – 比较 SageMaker 推理选项

在下图中,我们整理了在选择推理工作负载实现时,您需要注意的几个决策点:

图 10.8 – 推理选项的选择算法

图 10.8 – 推理选项的选择算法

请注意,您的工作负载配置不是静态的。以下是一些可能需要您重新考虑工作负载配置选择的非全面示例:

  • 流量模式的变化可能会导致扩展策略的变化

  • 用户服务水平协议(SLA)的变化可能会导致所选计算实例的变化和/或扩展策略的更新

  • 新版本的模型架构和/或可用的计算实例可能需要与基准进行基准测试,以衡量潜在的准确性或性能提升

因此,您应该计划并预算用于持续监控和工作负载优化的费用,作为初始工作负载设计的一部分。

使用 SageMaker 推理推荐器

选择最佳的推理配置需要相当多的工程和测试工作。为简化这一过程,AWS 最近推出了 SageMaker 推理推荐器,它为您提供了一种简便的方式,评估不同配置下实时端点的推理性能和成本。

推理推荐器将您的模型部署到具有不同配置的实时端点,进行负载测试,然后提供延迟和吞吐量度量,以及相关成本。根据生成的度量,您可以根据您的 SLA 和成本预算选择最合适的配置。SageMaker 推理推荐器提供以下基准:

  • 端到端 模型延迟(毫秒级)

  • 每分钟最大调用次数

  • 每小时成本每次推理成本

SageMaker 推理推荐器非常适合以下使用案例:

  • 寻找最佳实例类型。请注意,您可以提供自己感兴趣的实例类型列表进行基准测试,或者让 SageMaker 在所有支持的实例上对这个列表进行基准测试。

  • 由 SageMaker Neo 模型编制的基准。在这里,您可以比较原始模型与编译后的模型变体的性能。

  • 运行自定义负载测试。推理推荐器支持模拟不同的流量模式,以在不同条件下基准测试你的端点性能。因此,你可以使用 SageMaker 推理推荐器来基准测试和微调你的模型服务器配置、不同的模型版本等。

请注意,在撰写时,推理推荐器仅支持实时端点。因此,如果你需要基准测试不同的推理选项(例如,无服务器推理),你可能需要使用自定义基准测试和负载测试功能。此外,推理推荐器提供的基准统计信息以及支持的流量模式是有限的。

总结

在本章中,我们讨论了如何使推理工作负载实现和优化。我们介绍了亚马逊 SageMaker 提供的各种推理选项和模型托管选项,例如多模型、多容器和无服务器推理。然后,我们回顾了如何使用生产变体功能来推广和测试模型候选。

之后,我们提供了使用 SageMaker 部署保护措施的高级模型部署策略概述,以及使用亚马逊 CloudWatch 服务和 SageMaker 的模型监控功能进行工作负载监控。最后,我们总结了在定义推理工作负载配置时应使用的关键选择标准和算法。

posted @ 2025-07-12 11:41  绝不原创的飞龙  阅读(90)  评论(0)    收藏  举报