步骤 1:使用SageMaker分布式模型并行库修改自己的训练脚本 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

步骤 1:使用SageMaker分布式模型并行库修改自己的训练脚本

使用本节学习如何自定义训练脚本以使用 Amazon SageMaker 模型并行库的核心功能。要使用库特定的 API 函数和参数,我们建议您将本文档与 SageMakerPython SDK 文档中的SageMaker模型并行库 API 一起使用。

这些部分中提供的训练脚本示例经过简化,旨在突出显示使用该库时必须进行的更改。有关演示如何在SageMaker模型并行库中使用TensorFlow或PyTorch训练脚本的端到端、可运行的笔记本示例,请参阅。亚马逊 SageMaker 分布式训练笔记本示例

使用模型并行度库拆分训练脚本的SageMaker模型

修改训练脚本以设置模型拆分的方法有两种:自动拆分或手动拆分。

自动模型拆分

使用SageMaker模型并行库时,可以利用自动模型拆分,也称为自动模型分区。该库使用分区算法来平衡内存,最大限度地减少设备之间的通信并优化性能。您可以配置自动分区算法以优化速度或内存。

或者,您可以使用手动模型拆分。除非您非常熟悉模型架构并且对如何有效地分割模型有很好的了解,否则我们建议您自动拆分模型。

工作原理

自动分区发生在第一个训练步骤中,即首次调用 smp.step-deardeard 函数时。在此调用中,库首先在 CPU RAM 上构造模型的版本(以避免 GPU 内存限制),然后分析模型图并做出分区决策。基于此决定,将每个模型分区加载到 GPU 上,然后才执行第一步。由于这些分析和分区步骤,第一个训练步骤可能需要更长的时间。

在任一框架中,库都通过自己的后端管理设备之间的通信,后端针对Amazon基础架构进行了优化。

自动分区设计适应框架的特征,库在每个框架中进行粒度级别的分区,这在每个框架中更自然。例如,在中TensorFlow,可以将每个特定操作分配给不同的设备,而在中PyTorch,分配是在模块级别完成的,其中每个模块由多个操作组成。以下部分回顾了每个框架中的设计细节。

在第一个训练步骤中,模型并行库在内部运行跟踪步骤,该步骤旨在构造模型图并确定张量和参数形状。在此跟踪步骤之后,库将构造一棵树,该树由模型中的嵌套nn.Module对象以及从跟踪中收集的其他数据组成,例如存储nn.Parameters量和每个nn.Module对象的执行时间。

接下来,库从根目录遍历这棵树,运行分区算法,nn.Module将每个分区算法分配给一个设备,从而平衡计算负载(由模块执行时间衡量)和内存使用量(由总存储nn.Parameter大小和激活次数衡量)。如果多个模块nn.Modules共享相同的参数nn.Parameter,则将这些模块放置在同一设备上,以避免维护同一参数的多个版本。做出分区决定后,分配的模块和权重将加载到其设备上。

有关如何将smp.step装饰器注册到PyTorch训练脚本的说明,请参阅自动拆分为 PyTorch

模型并行库分析可训练变量的大小和图结构,并在内部使用图分区算法。该算法为每项操作提供设备分配,目标是最大限度地减少设备之间所需的通信量,但有两个限制:

  • 平衡每台设备中存储的变量数量

  • 平衡每台设备中执行的操作数量

如果您指定 speedoptimize(在 Python SDK 的模型并行度参数中),则该库会尝试平衡每台设备中的操作数量和tf.Variable对象数量。否则,它会尝试平衡的总大小tf.Variables

做出分区决定后,库会创建每个设备需要执行的子图的序列化表示形式,并将它们导入到每个设备上。分区时,库会将消耗相同内容的操作tf.Variable和属于同一 Keras 层的操作放置在同一设备上。它还尊重由TensorFlow施加的主机托管限制。这意味着,例如,如果有两个 Keras 层共享一个tf.Variable,则属于这些层的所有操作都将放在单个设备上。

有关如何将smp.step装饰器注册到PyTorch训练脚本的说明,请参阅自动拆分为 TensorFlow

框架间自动模型拆分的比较

在中TensorFlow,基本计算单位是 atf.Operation,并将模型TensorFlow表示为 tf.Operation s 的有向无环图 (DAG),因此模型并行度库对此 DAG 进行分区,以便每个节点连接到一台设备。至关重要的是,tf.Operation对象足够丰富,具有可自定义的属性,从某种意义上说,它们是普遍的,因为每个模型都保证由此类对象的图形组成。

PyTorch另一方面, 它没有足够丰富和普遍的同等行动概念. 中PyTorch具有这些特征的最接近的计算单位是 nnn.Module,它的粒度级别要高得多,这就是库在中进行此级别分区的原因。PyTorch

手动拆分模型

如果您想手动指定如何跨设备对模型进行分区,请使用smp.partition上下文管理器。有关如何设置上下文管理器以进行手动分区的说明,请参阅以下页面。

要在修改后使用此选项,在步骤 2 中,您需要将设置为auto_partition,然后default_partition在 SageMaker Python SDK 的框架估算器类中定义。False任何未通过smp.partition上下文管理器明确放置在分区上的操作都将在上执行default_partition。在这种情况下,会绕过自动拆分逻辑,每个操作都是根据您的规范进行的。基于生成的图形结构,模型并行度库自动创建流水线执行计划。