SageMaker 分布式模型并行性最佳实践 - Amazon SageMaker
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),如所示。与针对 Amazon 基础架构进行了优化的SMDDP库的兼容性一个例外是 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 不会对激活进行分片,激活可能会导致 GPU 耗尽内存。在这种情况下,您可以通过使用 SMP v2 的张量并行性构成 PyTorch FSDP 来进行高效训练。

参考配置

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

模型 模型尺寸(模型参数的数量) 实例数 分片数据并行度 张量并行度 激活检查点 激活分载 批次大小
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 控制台和 Amazon 监控和记录训练作业 CloudWatch

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

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

  2. 选择训练作业

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

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

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

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

权限

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