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

模型管道传输

SageMaker 模型并行性库的核心功能之一是管道并行性，它决定了模型训练期间进行计算和跨设备处理数据的顺序。管道传输技术通过让 GPU 同时计算不同数据样本，实现模型并行性中真正的并行处理，并克服由于顺序计算造成的性能损失。使用管道并行性时，训练作业以管道方式按照微批次执行，以最大限度地提高 GPU 使用率。

管道执行计划

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

例如，根据管道计划和模型分区，GPU i 可能会对微批次 b 执行（向前或向后）计算，而 GPU i+1 则对微批次 b+1 执行计算，从而使两个 GPU 同时处于活动状态。在单次向前或向后传递期间，根据分区决策，单个微批次的执行流程可能会多次访问同一个设备。例如，位于模型开头的操作，与位于模型末尾的操作可能会放在同一设备上，而两者之间的操作位于不同的设备上，这意味着此设备被访问了两次。

该库提供了简单和交错这两种不同的管道计划，在 SageMaker Python SDK 中可以使用 pipeline 参数进行配置。在大多数情况下，交错管道可以通过更有效地利用 GPU 来实现更好的性能。

交错管道

在交错管道中，尽可能优先考虑向后执行微批次。这样可以更快地释放内存用于激活，从而更有效地使用内存。这还允许将微批次数量扩展到更大，从而减少 GPU 的空闲时间。在稳态下，每个设备在向前传递和向后传递之间交替运行。这意味着一个微批次的向后传递可能会在另一个微批次的向前传递完成之前运行。

上图说明了 2 个 GPU 上交错管道的执行计划示例。在图中，F0 表示微批次 0 的向前传递，B1 表示微批次 1 的向后传递。更新表示优化器对参数的更新。GPU0 始终尽可能优先考虑向后传递（例如，在 F2 之前执行 B0），这样可以清除在之前激活中使用的内存。

简单管道

相比之下，简单管道则为每个微批次运行向前传递，然后再开始向后传递。这意味着它仅对自身中的向前传递和向后传递阶段进行管道传输。下图说明了这种方式在 2 个 GPU 上的工作原理。

特定框架中的管道执行

使用以下部分来了解针对 TensorFlow 和 PyTorch， SageMaker 的模型并行库做出的特定于框架的管道计划决策。

使用 TensorFlow 的管道执行

下图中的示例演示了使用自动模型拆分，按模型并行性库的 TensorFlow 图形分区。拆分图形时，生成的每个子图形都会被复制 B 次（变量除外），其中 B 是微批次的数量。在此图中，每个子图形被复制 2 次 (B=2)。在子图形的每个输入中插入一个 SMPInput 运算，在每个输出中插入一个 SMPOutput 运算。这些操作与库后端通信，以在彼此之间传输张量。

下图是 2 个子图形的示例，以 B=2 进行拆分，并添加了梯度运算。SMPInput 运算的梯度是 SMPOutput 运算，反之亦然。这使梯度能够在反向传播期间向后流动。

此 GIF 演示了一个示例交错管道执行计划，其中 B=2 个微批次和 2 个子图形。每个设备按顺序执行其中一个子图形副本，以提高 GPU 利用率。随着 B 增大，空闲时段的比例趋近于零。每当需要对特定的子图形副本执行（向前或向后）计算时，管道层都会向相应的蓝色 SMPInput 运算发出信号以开始执行。

计算出单个小批次中所有微批次的梯度后，该库就会合并各个库的梯度，然后将其应用于参数。

使用 PyTorch 的管道执行

从概念上而言，管道传输遵循了与 PyTorch 类似的概念。但是，由于 PyTorch 不涉及静态图形，因此模型并行库的 PyTorch 功能使用了更为动态的管道模式。

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

由于当前设备在远程子模块执行期间处于空闲状态，当前微批处理的本地执行暂停，库运行时系统会将执行切换到当前设备可以主动处理的另一个微批次。微批次的优先级由所选的管道计划决定。对于交错管道计划，只要有可能，就会优先考虑处于计算后退阶段的微批次。