Amazon Web Services 在简化机器学习模型上的探索

发布于: Jul 22, 2022

Amazon Web Services 致力于简化机器学习模型，为客户提供便利，他们希望获得一种更简单、更具成本效益的解决方案

Amazon Web Services 致力于简化机器学习模型，为客户提供便利，客户告诉我们，以往几种选项在处理超大型模型方面都不太理想。他们希望获得一种更简单、更具成本效益的解决方案。于是，我们给出了自己的答案。

推出 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
        },
    })