亚马逊AWS官方博客
Amazon SageMaker使用数十亿条参数简化深度学习模型训练
原文链接：
今天，我们高兴地宣布，Amazon SageMaker已经在大型深度学习模型的训练方面迎来简化，帮助更多缺少丰富硬件资源的客户踏入高阶深度学习之门。
过去十年以来，作为机器学习技术的重量级子集，深度学习（DL）席卷了整个世界。深度学习算法以神经网络为基础，具有非凡的效能，可以提取隐藏在大量非结构化数据（例如图像、视频、语音或文本）当中的信息模式。实际上，深度学习在各种以往只能由人类完成的高复杂度任务上，特别是计算机视觉与自然语言处理方面，已经迅速取得令人瞩目的成就。实际上，随着深度学习在ImageNet大规模视觉识别挑战赛（ILSVRC）、通用语言理解评估（GLUE）或者斯坦福问题解答数据集（SQUAD）等参考任务上不断提升自身能力，其创新速度也全面进入新的层级。
为了应对越来越复杂的任务，深度学习研究人员设计出愈发精妙的模型、添加更多神经元层与更多连接，旨在提高模式提取能力与预测准确性，而这一切又直接影响到模型的体量。例如，使用100 MB的ResNet-50模型已经能够在图像分类方面取得良好结果。但如果要实现对象检测或者实例区分等更为困难的任务，我们就得使用更大的模型，例如大小约为250 MB的Mask R-CNN或者YOLO v4。
毫无疑问，模型体量的增长也会影响到模型训练所需要的时间及硬件资源。也正因为如此，图形处理单元（GPU）才成为长期以来大型深度学习模型训练与调优方面的首选方案。为了适应GPU当中的大规模并行架构与庞大的内置内存，行业开发出分批训练技术。相关方案会向GPU一次性发送多个数据样本，借此减少通信开销，大大加快了训练速度。例如，Amazon Elastic Compute Cloud (EC2) p4家族提供的英伟达A100包含7000多个计算核心与40 GB的高速内存，看来足以支持用大量数据训练超大模型，是不是？
遗憾的是，没这么简单……OpenAI GPT-2拥有约15亿项参数，T5-3B拥有约30亿项参数，GPT-3甚至包含约1750亿项参数，这些专司自然语言处理的庞然大物需要占用几十甚至数百PB的GPU内存。同样的，处理高分辨率3D图像的最新模型同样远超GPU内存的容纳能力，甚至就连单一批次都很消化。
为了应对这样的严苛挑战，深度学习研究人员只能尝试将多种技术整合起来，例如：
- 购买更多强大的GPU，但对于某些特定模型，使用顶级GPU已经成为一种必选项。
- 转而使用功能较弱的模型，牺牲一定准确性。
- 建立梯度检查点，即将中间训练结果保存至磁盘上（而非将所有内容留存在内存内），但这会将训练速度降低20%至30%。
- 实现模型并行，即手动拆分模型，并在不同GPU上训练模型的某个部分。毫无疑问，即使对于专家来说，这也是一项极度困难、耗时且充满不确定性的任务。
客户告诉我们，以上几种选项在处理超大型模型方面都不太理想。他们希望获得一种更简单、更具成本效益的解决方案。于是，我们给出了自己的答案。
推出SageMaker全新模型并发库
SageMaker模型并行库能够在多个GPU之间对模型进行高效分区，借此消除准确性降低因素或者复杂的手动操作。此外，得益于这种横向扩展的模型训练方法，客户不仅能够摆脱内存瓶颈，处理大型模型，同时也能够使用更多单体配置较低、但更加经济高效的GPU。
作为一项新功能，模型并发库目前支持TensorFlow与PyTorch，而且客户只需要对代码稍加调整。启动训练作业时，您可以指定要针对速度、还是针对内存用量来优化模型。接下来，Amazon SageMaker会代替您运行初始分析作业，借此判断模型的计算与内存资源需求。接下来，结论将被馈送至分区算法处，此算法将决定如何拆分模型、以及如何将模型分区映射至各GPU，同时保证传输通信量最低。分区决策的结果将被保存在文件内，此文件则以输入的形式被传递至实际训练作业当中。
可以看到，SageMaker负责打理一切。如果需要，您也可以手动配置模型、进行分区，而后在SageMaker上进行训练。
在讨论代码之前，我们先简单了解一下其内部工作原理。
通过模型分区与微批次机制完成训练
由于运行在不同GPU上的模型分区需要彼此提供正向输入（激活值），因此跨分区处理小批次处理作业会导致特定分区长期处于繁忙状态，而其他分区则处于停滞状态。
为了避免这种低效问题，我们将小批次进一步拆分为多个微批次，确保每个微批次都能在不同GPU上并发处理。例如，GPU #1可以正向传播微批次n，而GPU #2可以对微批次n+1进行同样的正向传播。所有激活值皆可保存起来，随时准备被传递至下一分区。
对于反向传播，分区同样需要相互传递输入值（梯度）。由于各分区无法同时进行正向与反向传播，因此我们可以等待所有GPU在自己的微批次上完成正向传播，而后再运行相应的反向传播。Amazon SageMaker提供这种简单模式。
此外，SageMaker还提供另一种高效选项，即交错模式。在这里，SageMaker将根据微批次的数量复制分区。例如，在使用2个微批次时，每个GPU都将运行已接收分区的两个副本。每个副本分别与运行在其他GPU上的分区进行协作，借此实现正向或反向传播。
在此基础上，整套架构就相当于是由2个相同的分区处理4个不同的微批次。
综上所述，SageMaker就是使用这样的方式最大程度实现不同微批次之前的正向与反向交错。
现在，让我们看看如何将这套体系与TensorFlow共同使用。
在Amazon SageMaker中实现模型并发性
借助SageMaker模型并发库，大家可以轻松在自己的TensorFlow代码中实现模型并发（过程与PyTorch相似）。以下为具体方法：
- 定义并初始化分区配置。
- 使用标准Keras子类，让模型成为
DistributedModel类的一个子类。 - 使用
@smp.step编写并装饰训练函数，此函数代表模型的正向与反向step。此函数将根据上一节中描述的架构实现管道化处理。 - 作为可选项，您也可以在管道的评估函数中执行相同的操作。
下面，我们将调度包含4个英伟达V100 GPU的ml.p3.8xlarge实例，对MNIST数据集进行简单的卷积网络训练。
首先，我们初始化模型并发API。
import smdistributed.modelparallel.tensorflow as smp
smp.init()
之后，将DistributedModel子类化并构建模型。
class MyModel(smp.DistributedModel):
def __init__(self):
super(MyModel, self).__init__()
self.conv = Conv2D(32, 3, activation="relu")
self.flatten = Flatten()
self.dense1 = Dense(128)
self.dense2 = Dense(10)
. . .
以下为训练函数的内容。
def forward_backward(images, labels):
predictions = model(images, training=True)
loss = loss_obj(labels, predictions)
grads = optimizer.get_gradients(loss, model.trainable_variables)
return grads, loss
接下来，我们可以像往常一样使用SageMaker SDK中提供的TensorFlow估计器进行训练。这里只需要添加模型并发配置：2个分区（因此需要在2个GPU上进行训练）与2个微批次（因此每个分区对应2个副本），且启用交错模式。
smd_mp_estimator = TensorFlow(
entry_point="tf2.py",
role=role,
framework_version='2.3.1',
pv_version='py3',
instance_count=1,
instance_type='ml.p3.16xlarge',
distribution={
"smdistributed": {
"modelparallel": {
"enabled":True,
"parameters": {
"microbatches": 2,
"partitions": 2,
"pipeline": "interleaved",
"optimize": "memory",
"horovod": True,
}
}
},
"mpi": {
"enabled": True,
"processes_per_host": 2, # Pick your processes_per_host
"custom_mpi_options": mpioptions
},
}
)
立即体验
通过以上方式，模型并发库显著降低了大规模深度学习模型的训练门槛。现在，您已经可以在提供Amazon SageMaker服务的各个Amazon Web Services区域中使用这项服务，且无需承担任何额外成本。
您可以参考相关示例立即体验，我们也期待听取您关于使用体验的感受。欢迎大家通过Amazon Web Services客户支持渠道以及Amazon SageMaker论坛与我们交流。