分布式训练与高效扩展：使用 Amazon SageMaker 模型并行和数据并行库 机器学习博客

在 Amazon SageMaker 中分布式训练与高效扩展

关键要点

自从 2022 年 12 月 ChatGPT 发布以来，分布式深度学习领域在大型语言模型LLMs的训练上取得了巨大发展。当前 LLM 的参数数量已经从十亿增长到数万亿，并且通常无法在单个加速器设备如 GPU或单个节点如 mlp532xlarge上进行训练，归因于内存限制。因此，客户在训练 LLM 时往往需要将工作负载分配到数百甚至数千个 GPU 上。在如此大规模的分布式训练中，实现高效训练是一个突出的挑战。近年来，分布式训练社区引入了三维并行性数据并行、管道并行和张量并行及其他技术如序列并行和专家并行来应对这些挑战。

2023 年 12 月，亚马逊推出了 SageMaker 模型并行库 20 (SMP)，该版本在大模型训练中实现了最佳效率。同时也发布了 SageMaker 分布式数据并行库 (SMDDP)。此版本相较于 1x 进行了重要更新：SMP 现在与开源 PyTorch 的 全切分数据并行 (FSDP) API 集成，允许用户在训练大模型时使用熟悉的接口，并且兼容 Transformer Engine (TE)，首次解锁张量并行技术。

在本文中，我们探讨了 Amazon SageMaker包括 SMP 和 SMDDP所带来的性能优势，以及如何利用这些库在 SageMaker 上高效地训练大型模型。我们展示了 SageMaker 在 mlp4d24xlarge 集群上最多 128 个实例以及 FSDP 混合精度bfloat16下训练 Llama 2 模型的基准表现。首先，我们展示了 SageMaker 的近线性扩展效率，接着分析各项功能对优化吞吐量的贡献，最后探索在张量并行下以不同的序列长度最高达到 32768进行高效训练的方案。

SageMaker 的近线性扩展效率

为了减少 LLM 模型的整体训练时间，在向大集群千个 GPU扩展时保持高吞吐量是至关重要的，尤其是考虑到节点间的通信开销。在本节中，我们展示了在 p4d 实例上，当固定总问题规模时，通过改变 GPU 数量实现的稳健和近线性扩展效率，同时调用 SMP 和 SMDDP。

在本节中，我们将演示 SMP 的近线性扩展性能。我们使用固定的序列长度 4096 来训练不同大小的 Llama 2 模型7B、13B 和 70B 参数，使用 SMDDP 后端进行集体通信，启用 TE，总批大小为 400 万，节点数从 16 到 128。下表总结了我们的最佳配置和训练性能模型每秒 TFLOPs。

如上表所示，总体表现展示了在不同的 sharded data parallel (sdp) 和张量并行 (tp) 组合下的最佳吞吐量。以 Llama 2 模型大小为 7B，序列长度为 4096 时，对于 32、64 和 128 个节点，分别实现了 970、916 和 841 的扩展效率相对于 16 个节点。无论模型大小不同，这些扩展效率都表现出相对稳定，并且随着模型大小的增大轻微上升。

SMP 和 SMDDP 对于其他序列长度如 2048 和 8192也展示了类似的扩展效率。

SageMaker 模型并行库 20 性能：Llama 2 70B

随着 LLM 领域持续发展，模型大小的增长和性能更新频繁。在本节中，我们展示了在 SageMaker 上使用固定模型大小 70B，序列长度 4096 和总批大小 400 万的 Llama 2 模型训练性能。为了与前表中列出的全局最佳配置和吞吐量进行比较基于 SMDDP 后端，通常采用 FSDP 混合切片和TE，下表扩展了其他最佳吞吐量可能通过张量并行以及关于分布式后端NCCL 和 SMDDP、FSDP 切片策略全切片和混合切片、是否启用 TE 的额外规格默认。

最新的 SMP 和 SMDDP 版本支持多个功能，包括原生 PyTorch FSDP、扩展和更灵活的混合切片、与 TE 的集成、张量并行和优化的全聚合操作。为了更好理解 SageMaker 如何实现高效的 LLM 分布式训练，我们探讨了 SMDDP 的增量贡献与以下 SMP 核心特性：

FSDP 全切片：SMDDP 在 NCCL 上的优化

如前表所示，当模型使用 FSDP 进行完全切片时，32 或 64 个节点上 NCCLTFLOPs #0与 SMDDPTFLOPs #1的吞吐量相当，但在 128 个节点时，从 NCCL 到 SMDDP 的性能提升高达 504。

在较小的模型规模下，我们观察到 SMDDP 对 NCCL 的显著持续改进，且在较小的集群规模上就能有效减轻通信瓶颈。

FSDP 混合切片以降低通信成本

在 SMP 10 中，我们推出了 分片数据并行 这一由亚马逊内部 MiCS 技术驱动的分布式训练技术。在 SMP 20 中，我们引入了 SMP 混合切片，这是一种更灵活的技术，使得模型可以在一部分 GPU 上进行切片，而不是在所有训练 GPU 上正如 FSDP 全切片的情况。对于中等规模的模型，不需要在整个集群中切片，即可满足每个 GPU 的内存约束。这使得集群上可以有多个模型副本，同时每个 GPU 运行时与更少的对等方进行通信。

SMP 的混合切片能够有效地实现模型切片，范围覆盖从不发生内存溢出的问题到全集群规模即全切片。

以下图描述了在 tp = 1 的情况下通过 sdp 的吞吐量依赖性，为了简化展示。虽然这不一定与前表中 NCCL 或 SMDDP 全切片的最佳 tp 值相同，但数字非常接近。它清晰地验证了在大型集群规模128 个节点下，从全切片切换到混合切片的价值，这适用于 NCCL 和 SMDDP。在较小的模型规模时，混合切片的显著改进在较小的集群大小下就已开始体现，而随着集群规模的增加，差距仍在扩大。

TE 的性能提升

TE 旨在加速 NVIDIA GPU 上的 LLM 训练。尽管在 p4d 实例上不支持 FP8，但我们仍然看到了 TE 的显著加速。

基于使用 SMDDP 后端训练的 MiCS，TE 在所有集群规模上引入了一致的吞吐量提升唯一例外是 128 个节点上的全切片，即便在禁用张量并行的情况下张量并行度为 1。

在较小模型规模或不同序列长度条件下，TE 的提升稳定并且非常显著，提升范围约为 3 到 76。

在资源设置较低情况下的激活卸载

在资源设置较低的情况下节点数量较少，启用激活检查点后，FSDP 可能遭遇较高的内存压力在最坏的情况下甚至会出现内存溢出。对于这种受到内存瓶颈的场景，启用激活卸载可能是提升性能的一个选择。

例如，如前所述，尽管 Llama 2 在模型大小为 13B、序列长度为 4096 时，能够在 32 个节点的情况下进行最佳训练而无需激活卸载，但在限制为 16 个节点时，启用激活卸载可获得最佳吞吐量。

通过 SMP 张量并行实现长序列训练

长序列的需求在 LLM 社区中越来越受到重视，尤其适用于长时间对话和上下文。因此，我们在下表中报告了 Llama 2 的不同长序列吞吐量。表中展示了在 SageMaker 上，以各种长度从 2048 到 32768训练 Llama 2 的最佳吞吐量。对于序列长度 32768，使用 32 个节点和总批大小 4 万的情况下，原生 FSDP 训练是不可行的。

在集群规模较大且给定固定的总批大小的情况下，某些模型训练可能会因没有内置的管道或张量并行支持而无法通过原生 PyTorch FSDP 完成。根据上表，在总批大小为 400 万、32 个节点和序列长度为 32768 时，每个 GPU 的有效批大小为 05例如，tp = 2，批大小为 1，在没有张量并行的情况下这将是不可能的。

结论

在本文中，我们展示了在 p4d 实例上利用 SMP 和 SMDDP 实现高效 LLM 训练的过程中，多个关键特性对性能贡献的显著，例如 SMDDP 在 NCCL 的优化、灵活的 FSDP 混合切片技术、TE 的集成，以及启用张量并行以支持长序列长度。经过大量不同模型、模型规模和序列长度的设置测试，显示出强大的近线性扩展效率，最多可支持 128 个 p4d 实例。总的来说，SageMaker 依然是 LLM 研究人员和从业者的一项强大工具。

想了解更多信息，请参阅 SageMaker 模型并行库 v2，或通过 smmodelparallelfeedback@amazoncom 联系 SMP 团队。

感谢

我们感谢 Robert Van Dusen、Ben Snyder、Gautam Kumar 和 Luis Quintela 提出的建设性反馈和讨论。

作者介绍

Xinle Sheila Liu 是亚马逊 SageMaker 的软件开发工程师。她在闲暇时间喜欢阅读和户外运动。

Suhit Kodgule 是 AWS 人工智能组的 软件开发工程师，专注于深度学习框架。他在业余时间喜欢远足、旅行和烹饪。

Victor Zhu 是亚马逊网络服务的分布式深度学习软件工程师。他在旧金山湾区喜欢远足和玩棋类游戏。

Derya Cavdar 是 AWS 的软件工程师。她的兴趣包括深度学习和分布式训练优化。

Teng Xu 是 AWS AI 的分布式训练组软件开发工程师。他喜欢阅读。