SageMaker 分布式模型并行性最佳实践 - 亚马逊 SageMaker AI
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

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

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

使用 SageMaker 模型并行库 v2 (SMP v2) 运行分布式训练作业时,请遵循以下准则。

为分布式训练设置正确的配置

要估算并找到应用 SMP v2 提供的分布式训练技术的最佳起点,请查看以下列表。每个清单项目都讨论了使用 SageMaker 模型并行度库 v2 的核心功能 的优点和潜在的利弊得失。

配置提示

本节将介绍如何根据全局批次大小要求决定最佳模型配置,以获得最佳吞吐量。

首先,无论您的模型大小如何,我们都建议您采用以下设置。

  1. 使用您可以使用的功能最强大的实例类型。

  2. 始终开启混合精度,因为它能大大提高性能并减少内存。我们建议您使用 bfloat16,因为它比 float16 更精确。

  3. 只要适用,就打开SageMaker 分布式数据并行度库(而不是使用 NCCL),如所示。与针对基础架构进行了优化的 SMDDP 库的兼容性 Amazon一个例外是 tensor-parallelism-only用例(hybrid_shard_degree = 1tensor_paralle_degree > 1)。

  4. 如果您的模型有超过 600 亿个参数,我们建议您使用 延迟参数初始化。您还可以使用延迟参数初始化来加速任何模型的初始化。

  5. 我们建议您启用 激活检查点

根据您的模型大小,我们建议您先从以下指南开始。

  1. 使用分片数据并行。

    1. 根据您打算在 GPU 内存中容纳的批次大小,选择适当的分片数据并行度。通常情况下,您应该从最低度开始,以使模型适合 GPU 内存,同时尽量减少网络通信的开销。如果您看到缓存正在刷新的警告,我们建议您提高分片程度。

    2. 根据最大本地批次大小和所需的全局批次大小(如有)确定 world_size

    3. 您可以尝试激活卸载。根据不同的应用场景,它可以满足您的内存需求,而无需提高分片程度,这意味着通信量更少。

  2. 同时使用 FSDP 的分片数据并行性和 PyTorch SMP v2 的张量并行性,如中所述。张量并行性

    1. 在大型集群上进行训练时,如果仅使用 FSDP,全局批次大小可能会变得过大,从而导致模型收敛问题。通常情况下,大多数研究工作都会将令牌批次控制在 400 万枚以下。在这种情况下,您可以通过将 PyTorch FSDP 与 SMP v2 的张量并行性合成 FSDP 来减少批量大小来解决问题。

      例如,如果您有 256 个节点,序列长度为 4096,即使每个 GPU 的批次大小为 1,也会导致全局批次大小为 800 万个令牌。但是,当您使用张量并行度为 2 且每个张量并行组批次大小为 1 时,则每个 GPU 的批次大小就变成了 1/2,也就是 400 万个令牌。

    2. 当使用较长的上下文长度(例如 8k)进行训练时,16k 的激活内存可能会变得非常大。FSDP 不会对激活进行分片,激活可能会 GPUs 导致内存不足。在这种情况下,您可以通过使用 SMP v2 的张量并行性构成 PyTorch FSDP 来进行高效训练。

参考配置

SageMaker 模型并行度训练小组根据实验提供以下参考点,使用Llama 2模型转换为SMP变换器模型torch.sagemaker.transform,并在序列长度为4096和混合精度(或)的ml.p4d.24xlarge实例上进行训练。FP16 BF16

模型 模型大小(模型参数的数量) 实例数 分片数据并行度 张量并行度 激活检查点 激活分载 批次大小
Llama 2 7B 1 8 1 TRUE FALSE 4
70B 32 256 1 TRUE FALSE 2
175B 64 128 4 TRUE TRUE 6

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

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

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

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

  2. 选择训练作业

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

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

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

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

权限

要使用模型并行度运行 SageMaker 训练作业,请确保您的 IAM 角色拥有正确的权限,例如: