SageMaker 模型并行库的核心功能 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅中国的 Amazon Web Services 服务入门

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

SageMaker 模型并行库的核心功能

Amazon SageMaker 的模型并行库通过提供自动模型拆分和复杂的管道执行计划,使模型并行性更易于访问。模型拆分算法可以针对速度或内存消耗进行优化。该库还支持手动分区。当您使用库时,训练将以流水线方式通过微批量执行,以最大限度地提高 GPU 的使用率。

在创建训练脚本并定义 SageMaker 时,可以使用几行代码来配置这些功能PyTorch要么TensorFlow估计器。有关库的这些核心功能的更多信息,请参阅以下部分。

注意

SageMaker 分布式培训库只能通过AmazonSageMaker 培训平台中的 TensorFlow、PyTorch 和 HuggingFace 框架的深度学习容器。要使用这些库,您必须使用 SageMaker Python SDK 或者通过适用 SDK for Python (Boto3) 使用 SageMaker API 或Amazon Command Line Interface. 在整个文档中,说明和示例都重点介绍了如何将分布式培训库与 SageMaker Python SDK 结合使用。

自动拆分模型

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

此外,您也可以使用手动拆分模型。我们建议自动分割模型,除非您非常熟悉模型架构并且对如何有效地对模型进行分区有一个很好的想法。

工作方式

自动分区是在第一个训练步骤中进行的,当时smp.step-装饰功能首先被调用。在这次调用过程中,库首先在 CPU RAM 上构建模型的版本(以避免 GPU 内存限制),然后分析模型图并做出分区决策。根据这一决定,每个模型分区都加载到 GPU 上,然后才执行第一步。由于这些分析和分区步骤,第一个训练步骤可能需要更长的时间。

在任何一个框架中,库都通过自己的后端管理设备之间的通信,后者经过优化Amazon基础设施。

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

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

接下来,库从根目录遍历这个树并运行分区算法,分配每个nn.Module用于平衡计算负载(通过模块执行时间衡量)和内存使用(以存储总量衡量)的设备nn.Parameter大小和激活)。如果多个nn.Modules分享同样的nn.Parameter,然后将这些模块放置在同一设备上,以避免维护同一参数的多个版本。一旦做出分区决策,分配的模块和权重就会加载到他们的设备上。

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

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

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

如果你指定speed为了optimize(在 Python SDK 的模型并行参数中),库试图平衡操作数量和tf.Variable每台设备中的对象。否则,它会尝试平衡tf.Variables.

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

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

在 TensorFlow 中,基本的计算单元是tf.Operation,TensorFlow 将模型表示为的有向非循环图 (DAG)tf.Operations,因此建模并行库分区这个 DAG,以便每个节点都转到一个设备。至关重要的是,tf.Operation对象具有足够丰富的可自定义属性,从这个意义上讲,它们是普遍的,因为每个模型都保证都包含这些对象的图形。

另一方面,PyTorch 没有足够丰富和普遍的等效运营概念。PyTorch 中具有这些特征的最接近的计算单位是nn.Module,粒度级别要高得多,这就是为什么库在 PyTorch 中确实在这个级别进行分区的原因。

手动拆分模型

如果要手动指定模型在不同设备之间的分区方式,则可以使用手动拆分模型smp.partition上下文管理器。

要使用此选项,请设置auto_partitionFalse,然后定义default_partition在 SageMaker Python 开发工具包中。任何未通过smp.partition上下文管理器在default_partition. 在这种情况下,会绕过自动拆分逻辑,每个操作都会根据您的规格进行。基于生成的图形结构,模型并行库自动创建流水线执行计划。

管道执行计划

SageMaker 的分布式模型并行库的核心功能是管道执行,它决定了模型训练期间进行计算和跨设备处理数据的顺序。Pipelining 是一种技术,通过让 GPU 在不同的数据样本上同时计算,实现模型并行性能的真正并行化,并克服由于顺序计算而导致的性能损失。

Pipelining 的基础是将小批次拆分成微批次,这些微批次逐个输入到训练管道中,并遵循库运行时间定义的执行计划。一个微批是给定训练小批次中的一个较小的子集。管道计划决定了每个时段由哪个设备执行哪个微批处理。

例如,根据管道计划和模型分区,GPUi可以在微批处理上执行(向前或向后)计算(向前b而 GPUi+1在微批处理上执行计算b+1,从而使两个 GPU 同时保持活动状态。在单次向前或向后传递期间,单个微批处理的执行流可能会多次访问同一台设备,具体取决于分区决策。例如,模型开头的操作可能与模型末尾的操作放在同一设备上,而两者之间的操作在不同的设备上,这意味着该设备被访问两次。

图书馆提供两种不同的管道时间表,简单的交错,可以使用pipeline参数在 SageMaker Python 开发工具包中。在大多数情况下,交错管道可以通过更有效地利用 GPU 来实现更好的性能。

交错管道

在交错管道中,尽可能优先向后执行微批量。这样可以更快地释放用于激活的内存,更有效地使用内存。它还允许扩展更多微批次的数量,从而减少 GPU 的空闲时间。处于稳定状态时,每台设备在向前和向后运行之间交替运行。这意味着一个微批次的向后传递可能会在另一个微批次的向前通过结束之前运行。

上图说明了超过 2 个 GPU 的交错管道的示例执行时间表。在图中,F0 表示微批次 0 的正向传递,B1 表示微批次 1 的向后传递。更新表示参数的优化程序更新。GPU0 总是尽可能优先向后传递(例如,在 F2 之前执行 B0),这允许更早清除用于激活的内存。

简单管道

相比之下,一条简单的管道在开始向后传递之前完成了每个微批次的正向传递。这意味着它只管道自己内部的前向传递和向后通过阶段。下图说明了超过 2 个 GPU 的工作原理的示例。

在特定框架中管道执行

使用以下章节了解 SageMaker 的分布式模型并行库为 Tensorflow 和 PyTorch 制定的特定于框架的管道调度决策。

使用 TensorFlow 执行管道执行

下图是一个由模型并行库使用自动模型拆分进行分区的 TensorFlow 图的示例。分割图表时,每个生成的子图将被复制 B 次(变量除外),其中 B 是微批次的数量。在此图中,每个子图复制 2 次(B=2)。网络 ACL 和安全组都允许 (因此可到达您的实例) 的发起 ping 的SMPInput将在子图的每个输入中插入操作,并且SMPOutput将在每个输出处插入操作。这些操作与库后端进行通信,以便彼此之间传输张量。

下图是用 B=2 拆分并添加渐变运算的 2 个子图示例。一个的渐变SMPInputOP 是SMPOutputop,反之亦然。这使渐变能够在反向传播期间向后流动。

本 GIF 演示了一个与 B=2 个微批和 2 个子图交错管道执行计划的示例。每台设备按顺序执行其中一个子图副本以提高 GPU 利用率。随着 B 的增长,空闲时段的比例将变为零。无论何时是对特定子图副本进行(向前或向后)计算的时候,管道层都会向相应的蓝色发出信号SMPInput要开始执行的操作。

一旦计算了来自单个小批次中所有微批次的渐变,库就会组合跨微批次的渐变,然后可以将渐变应用于参数。

使用 PyTorch ore 执行管道执行

从概念上讲,流水线在 PyTorch 中遵循了类似的想法。但是,由于 PyTorch 不涉及静态图形,因此模型并行库的 PyTorch 功能使用了更动态的流水线模式。

与在 TensorFlow 中一样,每个批次被分成多个微批次,在每台设备上一次执行一个微批次。但是,执行计划是通过在每台设备上启动的执行服务器来处理的。只要当前设备上需要放置在另一个设备上的子模块的输出,就会向远程设备的执行服务器发送执行请求以及输入张量发送到子模块。然后,服务器使用给定的输入执行此模块,并返回对当前设备的响应。

由于当前设备在远程子模块执行期间处于空闲状态,因此当前微批处理的本地执行暂停,库运行时将执行切换到当前设备可以主动处理的另一个微批处理。微批量的优先级取决于所选的管道时间表。对于交错管道计划,尽可能优先考虑处于计算后面阶段的微批量。