acc-dl-wkld-amz-sgmk-merge-1

亚马逊 SageMaker 深度学习负载加速指南（二）

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

译者：飞龙

协议：CC BY-NC-SA 4.0

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

训练一个大型的 深度学习（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	高达 10
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	40
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.8xlarge或p3.16xlarge，它们是多节点实例。在这种情况下，你将把模型的部分存储在多个 GPU 中。NVLink 互联在前向和反向传播过程中提供 GPU 之间的高速数据交换。

p3.8xlarge和p3.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 通信算法

• 其次是模型的大小及其精度（在前面的公式中为 D）。

例如，如果我们使用 Ring-AllReduce 算法来训练 BERT 模型（该模型包含大约 3.4 亿个参数）并采用半精度，每个训练设备在单次迭代中将发送和接收大约 650 MB 的数据。通信需要快速进行。单个设备的性能下降将导致整体训练过程的放缓。

引入 EFA

Amazon EFA 是一种网络设备，提供比传统 TCP 传输更低且更稳定的延迟。EFA 专门为高性能和机器学习应用场景设计，其中实例间通信对分布式任务至关重要。

EFA 提供了以下优势：

• 操作系统绕过功能，允许深度学习应用直接与网络接口硬件通信，从而提供低延迟和可靠的传输功能。

• 支持高性能消息协议，如MPI和NCCL。对于深度学习用例，我们特别关注 NVIDIA 的 NCCL 库，它为 GPU 设备提供高性能的通信例程。

使用 EFA 可以显著提高训练任务的性能。根据 AWS 基准测试，使用 EFA 可以使 BERT 在 ml.p4dl.24xlarge 的 32 个实例上训练速度比默认的 弹性网络适配器 (ENA) 快 130%。

在 SageMaker 上使用 EFA 不会产生额外费用。EFA 可用于 ml.p3dn.24xlarge、ml.p4d.24xlarge 和 ml.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 NCCL、Gloo和消息传递接口（MPI）。

• allreduce 操作用于聚合和平均张量或广播，将张量从一个训练进程发送到集群中的其他进程。通常，通信后端提供集体操作的实现。

现在我们了解了分布式训练的基本术语，让我们深入回顾一下数据并行。

数据并行分布式训练在你想要减少在多个训练设备上训练模型的时间时非常有用。每个训练进程都有一个全局模型的副本，但它在与其他进程并行的过程中，在独特的数据切片上进行训练（因此称为数据并行性）。在训练步骤结束时，每个训练进程与其他进程交换学习到的梯度更新。然后，梯度更新被平均并分发回所有训练进程，以便它们可以更新各自的模型副本。图 6.1展示了数据批次在一个数据并行的双节点双 GPU 集群中的分布情况：

图 6.1 – 数据并行概览

在设计数据并行训练任务时，你需要注意几个关键设计选择，以便调试和优化你的训练任务，诸如以下几点：

• 各个进程之间如何进行协调

• 各个计算进程如何相互通信

• 计算进程如何在训练集群中分布

在接下来的章节中，我们将讨论这些设计选项。

协调模式 – 参数服务器与 Allreduce

在分布式集群中协调计算进程有两种方式：使用专用的集中式协调器，以及使用对等协调，其中每个节点直接与集群中的一个或多个对等节点通信。在数据并行训练的背景下，集中式协调模式被称为参数服务器，其中参数服务器进程协调梯度更新的分发并维护全局模型副本。而对等模式被称为Allreduce，它是一种通过对等算法在训练进程之间分发梯度更新的方式。在图 6.2中，你可以看到这两种协调模式的区别：

图 6.2 – 参数服务器 (A) 和 Allreduce (B) 协调模式

参数服务器负责协调集群中的训练过程，具体包括以下内容：

• 为每个训练过程分配一组独特的数据记录

• 从每个训练进程接收梯度

• 聚合梯度并相应地更新模型权重

• 将更新后的模型发送回训练进程

参数服务器存储模型权重的主副本。对于较大的深度学习模型，可能无法将完整的模型存储在参数服务器上。此外，参数服务器可能会成为网络和计算瓶颈。在这种情况下，你可能需要引入多个参数服务器，它们将存储模型参数的子集，从而减少网络和内存的需求。多个参数服务器允许你扩展分布式训练以处理大型模型；然而，当在训练进程和参数服务器之间协调模型更新时，它会引入额外的复杂性，并可能导致网络拥塞。找到训练进程与参数服务器之间的最佳配置可能是一个艰巨的任务，需要经过相当多的反复试验才能找到最优配置。

Allreduce 算法采用点对点通信，每个训练进程仅与两个邻居交换梯度更新。一个秩为 i 的训练进程为独特的数据微批次计算梯度，从进程 i-1 接收梯度，并将接收到的梯度与自己计算的梯度汇总，然后将聚合后的梯度发送给节点 i+1。总的来说，每个进程将与它的同伴进行 次通信：

图 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 实例如p4d和p3dn提供了改进的网络能力来解决这个问题。

现在我们已经对数据并行性有了初步的直觉，接下来让我们积累一些实践经验，并为 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 引入了一个名为 allreduce 和 allgather 的概念，并将其组织成层次结构，从而实现更好的整体性能。此外，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

让我们突出几个启动脚本中的关键部分：

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.16xlarge、ml.p3dn.24xlarge和ml.p4d.24xlarge。SDDP 提供了一个与 Horovod 和 PyTorch DDP 非常相似的 API，这使得从开源实现切换到它变得非常容易。

SDDP 实现了一个修改版的 Allreduce 算法，并进行了多项优化，以提高整体训练性能，特别是在allreduce操作中的等待时间。如前所述，在同步 Allreduce 算法中，通常分布式的allreduce操作是瓶颈，并且随着训练集群规模的扩大，它变得更加低效。请查看图 6.6：

图 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.16xlarge、ml.p3dn.24xlarge和ml.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 – 简单模型并行和流水线模型并行

图表来源

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 – 结合流水线并行和数据并行

现在我们已经刷新了对关键模型并行方法的理解，接下来让我们回顾一下 SDMP 库——SageMaker 专有的模型并行实现。

使用 SDMP 进行工程训练

SDMP 是一个功能丰富的库，能够实现各种类型的模型并行性和混合并行性，并针对 SageMaker 基础设施进行了优化。它支持 TensorFlow 和 PyTorch 框架，并允许你在最少的代码更改下，自动将模型划分到不同的设备。像 SDDP 一样，SDMP 使用 MPI 在训练集群中协调任务，在 GPU 设备上进行前向和反向计算，在 CPU 设备上进行通信任务。

SDMP 具有一些显著特点，可以简化模型并行训练作业的开发，并优化训练时硬件的利用率：

• SDMP 支持任意的模型架构，并且对训练脚本的修改最小，不会产生任何精度上的惩罚。

• 自动化模型拆分将在训练集群中将模型划分到不同设备之间。你可以选择优化速度和内存利用率。此外，SDMP 还支持手动模型拆分（然而，在实践中，这种方法通常不是一个好的选择）。

• 交错管道是对简单模型管道的改进，允许你通过优先执行反向操作来最小化处理微批次的时间：

图 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 parallel 和 mpi。

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.Linear、nn.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 中随时间变化的准确率

• 在**图形（Graph）**标签页中，您可以看到模型的可视化表示，以及数据从输入到输出的流动方式。

• 0，表示我们的模型正在学习，因此绝对梯度值在下降：

图 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 标签页

• Operators 标签页让你了解特定操作符（如卷积或批量归一化）消耗了多少时间。在下图中，我们可以看到，例如，卷积层的反向传递占用了大部分 GPU 时间：

图 7.4 – TensorBoard Profiler 的 Operators 标签页

• Kernel 标签页显示了在特定 GPU 核心上执行的时间。比如在下图中，你可以看到各种 单精度通用矩阵乘法（SGEMM）内核占用了大部分时间：

图 7.5 – TensorBoard Profiler 的 Kernel 标签页

• Trace 标签页显示了分析的操作符和 GPU 内核的时间线，以及 CPU 和 GPU 设备之间的交接（例如，将数据输入从 CPU 转移到 GPU）：

图 7.6 – TensorBoard Profiler 的 Trace 标签页

• 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库（github.com/awslabs/sag… 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。由于我们将 model 和 criterion 对象都添加到 hook 中，因此我们应该期望看到模型参数（权重、偏置等），以及损失标量。

在我们的训练循环中，我们唯一需要修改的地方是通过在 smedebug.modes.Train 和 smedebug.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 – 训练和验证损失

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 对我们的训练代码进行性能分析了。你可以在 github.com/PacktPublis… 中的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 存储

• 在你的训练脚本中，你需要指定检查点处理程序（请参考深度学习框架的文档），并将模型检查点存储在指定的目录中——即 /opt/ml/checkpoints。如果 Spot 实例被中断，SageMaker 将自动将该目录的内容复制到 S3。Spot 实例恢复可用后，SageMaker 会从 S3 将检查点复制回 /opt/ml/checkpoints 目录：

图 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 能力

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 容量

TensorRT 在编译模型时提供了几种关键的优化机制（参见图 8.3）：

• 精度校准将权重和激活值转换为 INT8 精度类型，而不会影响准确性，从而最大化模型吞吐量。

• 层和张量融合将多个层和张量运算合并为单一计算，以优化内存利用率和延迟。

• 内核自动调优为给定的硬件加速器选择最佳的数据层和算法。

• 动态张量内存允许你高效地重用为张量分配的内存。

• 多流执行允许你并行处理多个输入。

这些优化大多数都会在没有用户输入的情况下自动发生。在编译时，你需要设置以下参数：

• 精度模式定义了模型参数将转换成的精度类型。TensorRT 允许你在几乎不影响准确性的情况下降低精度。较低的精度可以减少内存占用，从而加速内存绑定的操作。

• 输入批量大小设置单次推理请求中预期的样本输入数量。增大批量大小通常会提高整体系统吞吐量。然而，较大的批量大小需要更多的可用内存，并且可能增加推理请求的延迟。

• 最大内存大小定义了在推理时可用于模型的 GPU 内存量。

建议根据可用资源和延迟吞吐量服务水平协议（SLA），尝试各种这些参数的组合，以获得最佳性能。

根据 DL 框架模型，编译到 TensorRT 格式的路径不同。对于 TensorFlow，你可以使用 TensorFlow-TensorRT（TRT）集成库（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 模型：BERT，distilBERT，XLM-BERT，Robert，BioBERT，MarianMT，Pegasus，和 Bart

• 计算机视觉模型：Resnet，Renext，VGG，Yolo v3/v4/v5，SSD

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 – 使用单一内存检索的批量推理

动态批处理是 Neuron SDK 的一项功能，它允许您切分输入张量，使其匹配编译时使用的批量大小。请注意，动态批处理适用于若干合适的模型架构。

NeuronCore 流水线

每个 Inferentia 加速器由四个NeuronCores组成。流水线技术允许您将模型分片到多个 NeuronCore 上，将模型参数缓存到片上内存中。这使得您可以使用本地缓存的数据更快地处理网络运算符，并避免访问外部内存。根据 AWS 的说法，内部基准流水线通常使我们在没有批处理的情况下实现最高的硬件利用率。下图展示了流水线的示例：

图 8.5 – 将模型流水线化到三个 NeuronCores 中

在以下示例中，我们将在 AWS Inferentia 实例上编译并基准测试 ResNet50 模型。撰写本文时，Amazon SageMaker 不支持托管的笔记本实例。因此，我们使用了 Amazon EC2 inf1.xlarge实例，配置为8888。为此，您需要像这样设置实例的安全组：

图 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 概览

SageMaker Neo 能够显著减少额外的开发或设置工作。然而，它也有一些限制，这些限制可能适合或不适合你的具体使用案例。

• SageMaker Neo 主要支持计算机视觉模型，如图像分类、目标检测和语义分割。它不支持 NLP 模型架构等其他类型的模型。

• SageMaker Neo 支持多个深度学习框架，但它们通常落后几个主要版本。因此，如果你希望使用最新的模型架构和/或框架版本特性，你需要考虑其他编译选项（例如，使用 TensorRT 手动编译）。有关 SageMaker Neo 支持的最新详情，请参考 docs.aws.amazon.com/sagemaker/latest/dg/neo-supported-cloud.xhtml。

• SageMaker Neo 支持一系列云实例。截至目前，它支持 ml.c5、ml.c4、ml.m5、ml.m4、ml.p3、ml.p2 和 ml.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 推理服务上开发和部署它们的实际经验。