SageMaker 分布式模型并行性最佳实践 - 亚马逊 SageMaker AI
针对给定模型设置合适的配置修改训练脚本使用 SageMaker AI 控制台和 Amazon 监控和记录训练 Job CloudWatch权限

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

SageMaker 分布式模型并行性最佳实践

使用 SageMaker 模型 parallel 库运行分布式训练作业时，请遵循以下准则。

针对给定模型设置合适的配置

在纵向扩展模型时，我们建议您按顺序浏览以下列表。每个列表项都讨论了使用该库技术的优点以及可能出现的权衡取舍。

提示

如果使用库的部分功能可以很好地拟合模型，那么添加更多的模型并行度或者使用节省内存的功能通常并不能提高性能。

使用大型 GPU 实例类型

  • 在模型并行性领域，最好使用具有大型 GPU 内存的强大实例来处理模型并行操作产生的开销，例如将模型分成多个模型。 GPUs我们建议使用 ml.p4d 或 ml.p3dn 实例来训练大型 DL 模型。这些实例还配备了 Elastic Fabric Adapter (EFA)，可提供更高的网络带宽，并支持使用模型并行性进行大规模训练。

分片优化器状态

  • 分片优化器状态的影响取决于数据并行秩的数量。通常，较高的数据并行度（与计算节点的大小成正比）可以提高内存使用效率。

    当您想要缩小集群规模时，请确保检查优化器状态分片配置。例如，具有优化器状态分片的大型 DL 模型适合具有 16 的计算集群 GPUs （例如，两个 P4d 或 p4de 实例），可能并不总是适合具有 8 的节点 GPUs （例如，单个 P4d 或 p4de 实例）。这是因为 8 GPUs 的组合内存低于总内存 16，并且在 16 GPUs GPU 场景中分片时每个 GPU 所需的内存也高于在 16 GPU 场景中进行分片所需的每个 GPU 内存。 GPUs 因此，内存需求的增加使其可能不适合较小的集群。

    有关更多信息，请参阅 优化器状态分片。

激活检查点

  • 通过对一组模块使用激活检查点功能，可以提高内存效率。分组的模块越多，内存使用效率就越高。在对层的顺序模块执行检查点操作时，smp.set_activation_checkpointing 函数的 strategy 参数会将层分组在一起以进行检查点操作。例如，将两个或多个层分组在一起执行检查点操作，比逐个层执行检查点操作更节省内存，但代价是需要额外的计算时间来减少内存使用量。

    有关更多信息，请参阅 激活检查点。

张量并行性

  • 张量并行度应为二的乘方（2、4、8、...、2 n），其中最大度数必须等于每个节点的 GPUs 数量。例如，如果您使用带有 8 的节点 GPUs，则张量并行度的可能数字为 2、4 和 8。对于张量并行度，我们不建议使用随意的数字（例如 3、5、6 和 7）。当您使用多个节点时，张量并行度配置错误可能会导致需要跨节点来运行张量并行性；这会显著增加跨节点激活通信的开销，并会导致计算成本高昂。

    有关更多信息，请参阅 张量并行性。

跨节点的流水线并行性

  • 您可以在单个节点内以及跨多个节点来运行管道并行性。当您将管道并行性与张量并行性结合使用时，我们建议在多个节点上运行管道并行性，并将张量并行性保持在单个节点中。

  • 管道并行性带有以下三个旋钮：microbatches、active_microbatches 和 prescaled_batch。

    • 当您将张量并行性与管道并行性结合使用时，我们建议激活 prescaled_batch，这样可以增加每个模型并行组的批次大小，从而实现高效的管道传输。激活 prescaled_batch 后，在训练脚本中设置的批次大小，将变为针对每个没有 prescaled_batch 的秩设置的批次大小的 tp_size 倍。

    • 增加 microbatches 的数量有助于实现高效的管道和更好的性能。请注意，有效微批次大小等于批次大小除以微批次数量。如果您在保持批次大小不变的同时增加微批次数量，则每个微批次处理的样本数会减少。

    • active_microbatches 的数量是在管道处理期间同时处理的最大微批次数量。对于正在处理的每个活跃微批次，其激活和梯度都会占用 GPU 内存。因此，增加 active_microbatches 会占用更多的 GPU 内存。

  • 如果 GPU 和 GPU 内存都未得到充分利用，则在管道处理中增加 active_microbatches 可以实现更好的并行化。

  • 有关如何将张量并行性与管道并行性结合使用的更多信息，请参阅张量并行性与管道并行性相结合。

  • 要查找上述参数的描述，请参阅 SageMaker Python SDK 文档smdistributed中的参数。

将激活任务分载到 CPU

  • 确保此操作与激活检查点和管道并行性结合使用。为确保在后台执行分载和预加载，请在微批次参数中指定大于 1 的值。

  • 对激活进行分载时，您也许能够增加 active_microbatches，并且有时需要与微批次总数相匹配。这取决于对哪些模块执行了检查点操作以及模型是如何分区的。

    有关更多信息，请参阅 激活分载。

参考配置

SageMaker 模型并行度训练团队根据对 GPT-2 模型的实验，提供了以下参考点，序列长度为 512，词汇量为 50,000。

模型参数的数量 实例类型 管道并行性 张量并行性 优化器状态分片 激活检查点 预缩放批次 批次大小
100 亿 16 ml.p4d.24xlarge 1 4 True 每个转换器层 True batch_size=40
300 亿 16 ml.p4d.24xlarge 1 8 True 每个转换器层 True batch_size=32
600 亿 32 ml.p4d.24xlarge 2 8 True 每个转换器层 True batch_size=56, microbatches=4, active_microbatches=2

您可以根据上述配置进行推断，以估算模型配置的 GPU 内存使用量。例如，如果您增加了 100 亿参数模型的序列长度或将模型的大小增加到 200 亿，则可能需要先减小批次大小。如果模型仍然不合适，请尝试提高张量并行度。

修改训练脚本

  • 在训练脚本中使用 SageMaker 模型并行库的功能之前，请先查看 SageMaker 分布式模型并行库配置提示和陷阱。

  • 要更快地启动训练作业，请使用 SageMaker AI 本地模式。这可以帮助您在 SageMaker 笔记本实例上快速在本地运行训练作业。根据运行 SageMaker 笔记本实例的 ML 实例的规模，您可能需要通过更改模型配置（例如隐藏宽度、变压器层数和注意点）来调整模型的大小。在使用大型集群训练完整模型之前，请验证缩减模型在笔记本实例上是否运行良好。

使用 SageMaker AI 控制台和 Amazon 监控和记录训练 Job CloudWatch

要监控 CPU 内存利用率、GPU 内存利用率和 GPU 利用率等系统级指标，请使用通过 SageMaker AI 控制台提供的可视化效果。

  1. 在左侧导航窗格中，选择训练。

  2. 选择训练作业。

  3. 在主窗格中，选择要查看更多详细信息的训练作业名称。

  4. 浏览主窗格并查找监控部分以查看自动生成的可视化对象。

  5. 要查看训练作业日志，请在监控部分选择查看日志。您可以在中访问训练作业的分布式训练作业日志 CloudWatch。如果您启动了多节点分布式训练，则应该会看到多个带有 algo-n-1234567890 格式标签的日志流。algo-1 日志流跟踪来自主（第 0 个）节点的训练日志。

有关更多信息，请参阅 用于 CloudWatch 监控和分析训练作业的 Amazon 指标。

权限

要使用模型并行度或SageMaker 分布式训练示例笔记本运行训练作业，请确保您在 IAM 角色中拥有正确的权限，例如： SageMaker