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

# 步骤 1：使用 SageMaker分布式模型并行库修改自己的训练脚本
<a name="model-parallel-customize-training-script"></a>

使用本节学习如何自定义训练脚本以使用 Amazon A SageMaker I 模型并行度库的核心功能。要使用特定于库的 API 函数和参数，我们建议您将本文档与 Pyth *SageMaker on SDK* 文档 APIs中的[模型并SageMaker 行库](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)一起使用。

这些部分中提供的训练脚本示例经过简化，旨在重点介绍使用库时必须进行的更改。有关 end-to-end演示如何在 SageMaker 模型并行度库中使用 TensorFlow 或 PyTorch 训练脚本的可运行笔记本示例，请参阅。[亚马逊 SageMaker AI 模型并行度库 v2 示例](distributed-model-parallel-v2-examples.md)

**Topics**
+ [使用模型并行度库拆分训练脚本的 SageMaker 模型](#model-parallel-model-splitting-using-smp-lib)
+ [修改 TensorFlow 训练脚本](model-parallel-customize-training-script-tf.md)
+ [修改 PyTorch 训练脚本](model-parallel-customize-training-script-pt.md)

## 使用模型并行度库拆分训练脚本的 SageMaker 模型
<a name="model-parallel-model-splitting-using-smp-lib"></a>

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

### 自动模型拆分
<a name="model-parallel-automated-model-splitting"></a>

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

或者，您可以使用手动模型拆分。除非您非常熟悉模型架构，并且很好地掌握了如何有效地对模型进行分区的知识，否则我们建议自动拆分模型。

#### 工作原理
<a name="model-parallel-automated-model-splitting-how-it-works"></a>

自动分区发生在第一个训练步骤中，也就是首次调用 `smp.step` 修饰的函数时。在此调用期间，库首先在 CPU RAM 上构造模型的版本（以避免 GPU 内存限制），然后分析模型图形并做出分区决策。根据这个决策，每个模型分区加载到一个 GPU 上，然后才会执行第一步。由于这些分析和分区步骤，第一个训练步骤可能需要较长的时间。

在这两个框架中，库都通过自己的后端管理设备之间的通信，后端针对 AWS 基础架构进行了优化。

自动分区设计适应了框架的特性，库按照在每个框架中更自然的粒度级别进行分区。例如，在中 TensorFlow，可以将每个特定操作分配给不同的设备，而在模块级别中 PyTorch，分配是在模块级别完成的，其中每个模块由多个操作组成。以下部分介绍了每个框架中的具体设计。

##### 自动拆分模型 PyTorch
<a name="model-parallel-auto-model-split-pt"></a>

在第一个训练步骤中，模型并行性库在内部运行一个跟踪步骤，用于构造模型图形并确定张量和参数配置。在此跟踪步骤之后，库将构造一个树，该树包含模型中的嵌套 `nn.Module` 对象，以及通过跟踪收集的其他数据，例如所存储的 `nn.Parameters` 数量，以及每个 `nn.Module` 的执行时间。

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

有关如何将`smp.step`装饰器注册到 PyTorch 训练脚本的说明，请参阅[使用自动拆分 PyTorch](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16)。

##### 自动拆分模型 TensorFlow
<a name="model-parallel-auto-model-split-tf"></a>

模型并行性库分析可训练变量的大小和图形结构，并在内部使用图形分区算法。该算法为每个操作提供了设备分配，目的是最大限度地减少所需的设备间通信量，但要遵守两个限制：
+ 平衡每个设备中存储的变量数量
+ 平衡每个设备中执行的操作数量

如果您为 `optimize` 指定 `speed`（在 Python SDK 的模型并行性参数中），则库会尝试平衡每个设备中的操作数和 `tf.Variable` 对象数量。否则，它会尝试平衡 `tf.Variables` 的总大小。

做出分区决策后，库会创建每个设备需要执行的子图形的串行化表示形式，并将其导入到每个设备上。在分区时，库会将使用相同 `tf.Variable` 的操作，以及属于同一 Keras 层的操作放在同一个设备上。它还尊重由 TensorFlow施加的主机托管限制。举例而言，这意味着如果有两个 Keras 层共享一个 `tf.Variable`，则属于这些层的所有操作都将放在单个设备上。

有关如何将`smp.step`装饰器注册到 PyTorch 训练脚本的说明，请参阅[使用自动拆分 TensorFlow](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-23)。

##### 各个框架的自动模型拆分的比较
<a name="model-parallel-auto-model-split-comparison"></a>

在中 TensorFlow，计算的基本单位是 a`tf.Operation`，它将模型 TensorFlow 表示为 `tf.Operation` s 的有向无环图 (DAG)，因此，模型并行度库对此 DAG 进行了分区，使每个节点都转到一台设备。至关重要的是，`tf.Operation` 对象具有足够丰富的可自定义属性，而且它们是通用的，因为每个模型都确保由此类对象的图形组成。

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

### 手动模型拆分
<a name="model-parallel-manual-model-splitting"></a>

如果要手动指定如何跨设备对模型进行分区，请使用 `smp.partition` 上下文管理器。有关如何设置上下文管理器以进行手动分区的说明，请参阅以下页面。
+ [使用手动拆分 TensorFlow](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-manual)
+ [使用手动拆分 PyTorch](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16-hvd)

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