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

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

SageMaker 分布式并行模型的核心特征

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

当您创建训练脚本并定义 SageMaker PyTorch 或 Tensorflow 时,您可以使用几行代码来配置这些功能Estimator. 有关库这些核心功能的更多信息,请参阅以下部分。

注意

SageMaker 分布式培训库只能通过Amazon深度学习容器,用于 SageMaker 培训平台中的 TensorFlow、PyTorch 和 HugingFace 框架。要使用这些库,您必须使用 SageMaker Python 软件开发工具包或通过 SDK for Python (Boto3)(Bto3)或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.Operation,因此建模并行库分区此 DAG,以便每个节点都转到一个设备。至关重要的是,tf.Operation对象具有足够丰富的可定制属性,并且它们是通用的,因为每个模型都保证由这些对象的图形组成。

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

手动模型拆分

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

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

管道执行时间表

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

流水线是基于将小批次拆分为微批次,这些微批次逐个送入训练管道,并遵循库运行时定义的执行计划。A微粒是给定训练小批次的较小子集。管道计划确定每个时隙由哪个设备执行哪个微批处理。

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

该图书馆提供两种不同的管道时间表SIMPLE交错排序,可以使用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 进行分割,并添加了渐变运算。的渐变SMPInputop 是SMPOutputop,反之亦然。这使得渐变在反向传播过程中向后流动。

此 GIF 演示了一个带 B=2 微批次和 2 个子图的交错管道执行计划示例。每个设备按顺序执行其中一个子图复本,以提高 GPU 利用率。随着 B 的增大,空闲时隙的部分变为零。无论何时需要对特定的子图复本执行(向前或向后)计算,管线图层将信号到相应的蓝色SMPInput操作开始执行。

一旦计算了一个小批次中所有微粒的梯度,该库将跨微粒的梯度结合起来,然后将这些渐变应用于参数。

PyTorch 管道执行

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

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

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