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

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

SageMaker 分布式模型的核心功能并行

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

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

自动模型拆分

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

或者,您也可以使用手动模型拆分。我们建议自动进行模型拆分,除非您非常熟悉模型架构并很好地了解如何对您的模型进行高效分区。

如何使用

自动分区在第一个训练步骤(首次调用 smp.step装饰的函数)期间进行。在此调用期间,库首先在 CPU RAM 上构造模型版本(以避免 GPU 内存限制),然后分析模型图并做出分区决策。根据此决策,每个模型分区在 GPU 上加载,并且仅在此时执行第一步。由于这些分析和分区步骤,第一个训练步骤可能需要更长时间。

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

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

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

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

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

如果您speed为 指定 optimize (在 Python 开发工具包中的模型并行参数中),则库将尝试平衡每个设备中的操作和tf.Variable对象数。否则,它会尝试平衡 的总大小tf.Variables

做出分区决策后,该库会创建子图的序列化表示形式,每个设备需要执行该表示形式并将其导入到每个设备上。在分区时,库将使用属于同一 Keras 层一部分的相同 tf.Variable 和 操作的操作放置在同一设备上。它还遵循 TensorFlow施加的托管约束。这意味着,例如,如果两个 Keras 层共享一个tf.Variable,则属于这些层的所有操作都放置在单个设备上。

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

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

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

在 TensorFlow 中,基本的计算单位是 tf.Operation,TensorFlow 将模型表示为 tf.Operation的有向无环图 (DAG),因此,为此 DAG 建立并行库分区,以便每个节点转到一个设备。至关重要的是,tf.Operation对象足够丰富,具有可自定义的属性,并且它们是通用的,因为每个模型都保证包含此类对象的图形。

另一方面, PyTorch 没有足够丰富且通用的等效操作概念。具有这些特征的 PyTorch 中最近的计算单位是 nn.Module,它具有更高粒度级别,这就是该库在 PyTorch 中在此级别进行分区的原因。

手动模型拆分

如果要手动指定如何跨设备分区模型,则可以使用smp.partition上下文管理器进行手动模型拆分。

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

管道执行计划

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

管道传输功能基于将小批量拆分成多个小批量,这些小批量逐个注入到训练管道中,并遵循由库运行时定义的执行计划。微型批次是给定训练小批量的较小子集。管道计划确定哪个微批次由哪个设备针对每个时段执行。

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

该库提供两个不同的管道计划(简单交错),可使用 pipeline Python 开发工具包中的 SageMaker 参数配置这些计划。在大多数情况下,交错管道可以更高效地利用 GPUs 来实现更好的性能。

交错管道

在交错管道中,尽可能优先执行微批处理。这样可以更快地释放用于激活的内存,从而更高效地使用内存。它还允许将小批量的数量扩展得更高,从而减少 GPUs的空闲时间。在稳态时,每个设备在执行向前和向后传递之间交替。这意味着,一个微批处理的向后传递可在另一个微批处理的向前传递完成之前执行。

上图说明了交错管道 2 个 GPUs 的示例执行计划。在图中F0 表示微批次 0 的前向传递B1 表示微批次 1 的后向传递。Update (更新) 表示参数的优化程序更新。GPU0 始终尽可能优先向后传递(例如,在 F2 之前执行 B0这允许清除以前用于激活的内存。

简单管道

相比之下,简单管道在开始向后传递之前完成为每个微批次执行向前传递。这意味着,它只对前向传递和后向传递阶段进行管道。下图演示了这在 2 个 GPUs 中的工作原理的示例。

特定框架中的管道传输执行

可以使用以下部分了解 Tensorflow 和 PyTorch 的特定于框架SageMaker的管道计划决策的分布式模型并行库。

使用 TensorFlow的管道执行

下图是使用自动模型拆分按模型并行库分区的 TensorFlow 图表的示例。在拆分图表时,将复制 B 次每个生成的子图表(变量除外),其中 B 是微批处理的数量。在此图中,每个子图复制 2 次 (B=2)。在子图的每个输入中插入一个SMPInput操作,并在每个输出中插入一个SMPOutput操作。这些操作与库后端通信,以便相互传输张量。

下图是一个包含 2 个拆分的子图形的示例,其中 B=2 添加了梯度操作。SMPInput op 的梯度是一个 SMPOutput op,反之亦然。这使得梯度可在反向传播期间向后流动。

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

在计算单个小批量中所有微批处理的梯度后,库将跨微批量合并梯度,这些梯度随后可应用于参数。

使用 PyTorch的管道执行

从概念上说,管道传输遵循 PyTorch 中的类似概念。但是,由于 PyTorch 不涉及静态图形,因此模型并行库的 PyTorch 功能使用更动态的管道传输范例。

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

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