亚马逊 SageMaker 模型并行库中的新性能改进

基础模型是在海量数据上大规模训练的大型深度学习模型。可以进一步对其进行微调，以执行各种下游任务，并构成支持多个 AI 应用程序的核心支柱。最突出的类别是大语言模型 (LLM)，包括自动回归模型，例如经过训练以完成自然文本的 GPT 变体。LLM 通常包含数十亿个参数，因此它们很少能安装在单个加速器上，并且需要模型并行化技术。另一个类别是扩散模型，尤其是 稳定扩散 模型 ，它已将人工智能图像的生成推向了前所未有的里程碑，通过简单的文本描述可以生成出非凡的视觉效果。扩散模型通常比 LLM 小得多，分布式训练在促进开发方面仍然起着至关重要的作用。

SageMaker 模型并行 (SMP) 库 是一种大型模型训练解决方案，可在亚马逊 S ageMaker 平台上使用。 它可以与 PyTorch 模型集成，轻松应用一系列最先进的大型模型分布式训练技术进行大规模训练。今年早些时候，SMP推出了分 片数据并行性 ，这是一种分布式训练技术，由亚马逊内部的M IC s 技术提供支持。分片数据并行分片可以对数据并行工作线程中的参数、梯度和优化器状态进行建模。MIC 执行多项优化，包括标度感知分区，以提供近乎线性的可扩展性。在 “ 使用分片数据并行性 训练具有近线性缩放能力的巨型模型” 中，我们分享说 ，与序列长 度为 2048 的 30B 参数 GPT-2 模型上的 DeepSpeed Zero-3 相比，SMP 中的分片并行数据速度提高了 39.7%。

为了帮助我们的客户进一步最大限度地降低培训成本并缩短上市时间，我们很高兴在 SageMaker 并行模型中推出两项新的性能改进——S M DDP Collectives 和 FlashAttenction。 SMDDP Collections 是 亚马逊云科技 基础设施上性能最高的大型模型训练集体库，由 Sage Maker 分布式数据并行库提供。 FlashAttention 是由 Dao 等人介绍的。 ，它以 IO 感知方式重新实现注意力机制，降低了内存带宽需求并节省了注意力速度和内存占用。在 32 个 p4d.24xlarge 实例上训练 100B 参数 GPT-Neox 模型时，这两个组件共同推动我们的分片数据并行技术的速度提高了 30.58%。对于已经在支持的模型上并行使用分片数据的客户，无需更改任何代码即可受益于这些最新功能提供的性能提升。 Stabil ity AI是 稳定 扩散 模型 系列的发明者，该模型显示出无与伦比的图像生成能力，它选择使用SMP来构建基础模型。借助 SMP，Stability AI 在 32 个 p4d.24xlarge 实例上以 13B 参数的 GPT-Neox 实现了每个 GPU 163 Tflops，与 DeepSpeed 相比，速度提高了 58%。 你可以在 亚马逊云科技 re: Invent 2022 上对稳定人工智能首席执行官的演讲或这篇 博客文章中详细了解稳定人工智能的使命以及与 亚马逊云科技 的合作关系。

“我们在Stability AI的使命是奠定基础，通过人工智能激活人类的潜力。为了实现这一使命，我们需要在数百个加速计算实例上高效地训练开源基础模型。我们依靠 SageMaker 及其分布式训练库来优化性能并实施最先进的策略，在整个训练集群中分片模型和数据。这些优化降低了我们的培训成本，帮助我们更快地满足客户需求，并加快新模型的开发。”

— 稳定人工智能创始人兼首席执行官 埃马德·莫斯塔克。

在这篇博客文章中，我们将首先介绍 SageMaker 模型并行库中的最新性能改进。然后，我们将重新审视如何使用分片数据并行训练基础模型。最后，我们将对 13B、50B 和 100B 参数自动回归模型的性能进行基准测试，并总结未来的工作。

SageMaker 模型并行库中的新性能改进

从 亚马逊云科技 深度学习容器 (DLC) PyTorch 1.12.1 开始，SageMaker 模型并行库 v1.13 包含以下两个对提高训练性能至关重要的新组件。它们目前在启用了 弹性 结构适配器 (EFA) 的 ml.p4d.24xlarge 实例上可用：

1。亚马逊云科技 优化的 AllGather 来自 SMDDP 集体

在分片数据并行中，由于 GPU 上仅存在模型状态的分片，因此在向前或向后传递计算期间，需要一个 allGather 集合来收集来自分片组中所有 GPU 的完整参数集。在之前版本的 SageMaker 并行模型中，我们使用了 NVIDIA 集体通信库 (NCCL) 作为这些集体。但是，NCCL 是一个通用集体通信库，不是为 亚马逊云科技 基础设施设计的，即使启用了 EFA，也导致性能不理想。

之前，我们开发了 SMDDP Colle ctions 库，该库提供了 All-Reduce 集合体的优化实现，以提高纯数据并行训练的性能。为了使用分片数据并行性提高大型模型训练的性能，我们扩展了 SMDDP Collections 库，加入了 AllGather 集合体的优化实现。SMDDP Collections AllGather 的主要优势在于，它采用全对所有类型的通信模式进行节点间通信，使我们的集体具有高吞吐量并降低对延迟的敏感性。此外，我们的 AllGather 集合将与通信相关的处理转移到 CPU，从而腾出宝贵的 GPU 周期用于梯度计算，从而显著提高性能，尤其是在大型模型上。

2。闪光注意力

在现代变压器架构中，最大的内存消耗来源之一是自注意力层中的激活足迹。这是因为每个注意力头都会为每个输入计算 SxS 注意力矩阵，其中 S 是序列长度，并且该矩阵会经过多种操作，例如丢失、softmax 和矩阵乘法，每个中间输出都需要内存空间才能用于反向传播。

FlashAttent ion（Dao 等人 ）是斯坦福大学HazyResearch的最新创新，它以I/O感知方式重新实现了自我注意力机制。FlashAttenction 背后的主要见解是，自注意力机制受到进出 GPU 高带宽内存 (HBM) 的内存带宽的瓶颈。这意味着自注意力层可以在整个序列维度上按区块进行计算，每个区块一次都要经过整个自我注意力管道。区块的中间结果存储在高带宽 SRAM 中，从而避免了每次迭代向 HBM 的昂贵往返。尽管天真的实现会遇到softmax层的跨区块依赖问题，但FlashAttention引入了一个巧妙的实现来避开这种依赖关系。再加上向后传递中的重新计算，由于避免了 HBM 往返和存储 SxS 矩阵，FlashAttenction 可以节省大量内存并提高性能（Gpt-Neox 13B 的训练速度比 16 个 p4d 节点快 25%）。 你可以在 HazyResearch 的 F lashAttention 存储库中找到视觉效果和更多解释。

使用 SageMaker 并行模型大规模训练基础模型

要使用由 SMDDP Collections 提供支持的 SMP 训练基础模型，无需在分片数据并行训练作业中进行额外更改。如果你不熟悉并行使用分片数据，请关注 这 篇 完整的教程笔记本 和 博客文章 ，它将引导你完成从数据处理、定义和提交训练作业到监控训练日志的整个过程。 GPT-2 模型的即用型训练脚本可以在 train_gpt_simple.py 中找到。 要训练不同的模型类型，您可以按照 API 文档 了解 如何应用 SMP API。

我们重点介绍分片数据并行训练作业的 PyTorch Estimat o r 中的关键超参数，如下所示。 s m p_options 中的超参数 ddp_dist _backend 现在有一个新选项 “自动” 作为其默认值。 使用 “自动” 时，SMP 将使用经过 亚马逊云科技 优化的 AllGather 进行分片数据并行处理作业，否则将回退到 NCCL。您可以参考 本文档 了解支持的配置。 如果你想在 SMP 中并行运行分片数据，特别是使用 NCCL 作为首选的通信后端，你可以在 smp_options 中将 “ ddp_dist_ backend” 设置为 “nccl”。

import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "ddp_dist_backend": "auto", #OR "nccl" to disable SMDDP Collectives
        # To enable sharded data parallelism.
        # Here we shard model states across 128 GPUs.
        "sharded_data_parallel_degree": 128,  
    }
}

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    role=sagemaker.get_execution_role(),
    instance_type='ml.p4d.24xlarge',
    instance_count=32,
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        ...
    },
    ...
)

smp_estimator.fit(inputs=data_channels)

在最新的 smpv1.13 版本中，分片数据并行训练技术支持开箱即用的 BERT、RobertA、GPT-2、GPT-J、Gpt-neo 和 Gpt-Neox 等热门型号的 FlashAttenction。 这是通过在模型创建期间传递 tensor_parallelism=TRUE 来实现的，无需设置 tensor_parallel_degree 。 你可以在同一个训练脚本 train_gpt_simple.py 中找到一个示例 。

基准测试性能

我们在三种不同的模型尺度上对 SageMaker 模型并行库中的分片数据并行度进行了基准测试，以了解 FlashAttenction 和 亚马逊云科技 优化的 AllGather 这两项新功能如何促进性能改进。置放组无需在 SageMaker 上重现这些基准测试。

13B 参数 GPT-Neox

在此设置中，我们专注于理解 FlashAttenction 带来的性能提升，而我们将 亚马逊云科技 优化的 AllGather 排除在外。使用 FlashAttenction 可以节省大量的 GPU 内存，这有助于我们增加批量大小或降低分片程度，从而提高性能。如下结果所示，我们观察到，在16-64个p4d节点上的各种配置下，使用FlashAttenction的13B参数GPT-Neox模型的SMP平均加速约为20.4％。标准注意力计算期间的内存使用量以二次方式随着序列长度的增加而增加，但是 FlashAttenction 的内存使用量与序列长度成线性关系。因此，随着序列长度的增加以及使用更大的序列长度成为可能，FlashAttenction 会更加有用。 FlashAttenction在不牺牲模型质量的情况下提高了记忆效率，在过去的几个月中，包括与 Hugging Face Diff users和Mosaic ML的集成，在大型模型训练社区中迅速获得了关注。

50B 参数 Bloom

现在，我们来看看来自 SMDDP Collectives 的 亚马逊云科技 优化版 AllGather 如何使用 SMP 加速大型模型训练。我们对 50B 参数的 Bloom 模型进行了基准测试，比较了使用和不使用 亚马逊云科技 优化的 AllGather 集合体的性能。我们观察到，在 32 个节点到 64 个节点的训练作业中，SMDDP 集合可将模型训练速度提高多达 40%。由于可以更好地利用 p4d.24xlarge 实例可用的 400 Gbps 网络带宽，SMDDP 集合有助于提高性能。再加上将通信相关处理卸载到 CPU 的设计选择，有助于实现良好的计算到网络的重叠，从而优化性能。计算与网络的重叠在大型模型中尤其重要，因为节点间通信的数据大小会随着模型大小的增加而呈线性扩展。

100B 参数 GPT-Neox

最后，我们在启用了两个最新功能的情况下对 SMP 进行了基准测试。它表明，在参数为100B的GPT-Neox型号上，这个新版本的SMP v1.13比之前的版本快了30％。

对于未来的工作，我们将努力在 SMDDP 集体中支持 亚马逊云科技 优化的减少散射。Reduce-Scatter 集合对于对反向传递中计算的梯度求平均值和分片至关重要。我们预计这将在未来的版本中进一步加快 SMP 库的速度。

结论

在这篇文章中，我们讨论了 SageMaker 模型并行库中分片数据并行技术的两项最新性能改进。LLM 在提高机器学习模型的质量和可重用性方面表现出巨大的希望。亚马逊云科技 团队正在与客户紧密合作，不断降低培训成本和缩短上市时间。 你可以在 亚马逊 SageMaker 示例 GitHub 存储库中找到更多 SageMaker 模型并行示例 ，或者参加我们接下来的 分布式培训研讨会。 如果你有兴趣加快大型模型训练，请查看这些功能并告诉我们你建造了什么！

作者简介

Arjun Balasubramanian 是 亚马逊云科技 的高级软件工程师，专注于为分布式深度学习构建高性能、硬件加速的集体通信算法。他对大规模机器学习和网络系统广泛感兴趣。工作之余，他喜欢旅行和参加各种运动。

朱兆奇 是 亚马逊云科技 的软件开发工程师，专门研究分布式深度学习系统，并在 SageMaker 分布式数据并行库中工作。工作之余，赵奇对足球充满热情，希望在即将到来的赛季中不会获得任何红牌。

Can Karakus 是 亚马逊云科技 的高级应用科学家，负责在 亚马逊云科技 上优化大规模分布式深度学习。他的研究兴趣包括深度学习、分布式优化、分布式系统和信息论。工作之余，他喜欢骑自行车、旅行、阅读和学习。

Rahul Huilgol 是 亚马逊云科技 的高级软件工程师。他致力于开发分布式深度学习系统，旨在简化在云端训练大型深度学习模型并提高性能。在业余时间，他喜欢摄影、骑自行车和园艺。

Suhit Kodgule 是 亚马逊云科技 人工智能小组的一名软件开发工程师，致力于深度学习框架的开发。在业余时间，他喜欢远足、旅行和烹饪。

吴飞 是 亚马逊云科技 的软件工程师。他致力于云端大规模深度学习模型的分布式训练。工作之余，他喜欢篮球、游戏和烹饪。