该博客由Sarah Jane Hong CSO，Darryl Barnhart CTO和Latent Space首席执行官兼AWS的Prem Ranga首席执行官Ian Thompson共同撰写。
潜在空间是机器学习（ML）模型学习的抽象概念的隐藏表示。 例如，“狗”，“花”或“门”是潜在空间中的概念或位置。 在 潜在空间，我们正在开发一种引擎，使您可以通过语言和视觉提示来操纵和探索该空间。 潜在空间小组来自两个长期以来几乎没有重叠的领域：图形和自然语言处理（NLP）。 传统上，图像和文本的形式是分开处理的，每种形式都有其自身复杂，昂贵和脆弱的特征工程的历史。 NLP任务（例如文档理解或问题解答）通常与视觉任务（例如场景理解或渲染）几乎没有共同点，并且通常我们为每个任务使用非常不同的方法和模型。 但这正在迅速改变。
在单个共享潜在空间中的这些模式的合并解锁了从游戏到文档理解的新一代创意和商业应用程序。 但是，如理查德·萨顿（Richard Sutton）的“痛苦的一课”所强调的那样，在单一模型中解锁这些新应用程序会带来新的扩展挑战，以及过去几年中有关扩展定律的激动人心的工作。 为了使之成为可能，Latent Space正在开展前沿研究，以将这些模式融合到一个模型中，而且还要进行规模化和高效地进行。 这就是模型并行性出现的地方。
亚马逊SageMaker独特的自动模型划分和有效的流水线方法使我们能够以很少的工程工作量就可以采用模型并行性，并且我们将模型训练的规模扩展到超过1亿个参数（我们使用 p4d.24xlarge A100实例），这是我们的重要要求。 此外，我们观察到，当使用SageMaker模型并行库使用16个节点进行训练，使用38个GPU训练设置进行训练时，与以前的训练相比，我们记录了XNUMX％的效率提升。
培训大型变压器的挑战
在潜伏空间，我们正在将变压器模型中的语言和视觉与数十亿个参数融合在一起，以支持用户想象中的“分布”用例，或者用在现实世界中，而不会出现在训练数据中。 我们正在以两种不同的方式来应对扩展到数十亿个参数及更多参数所固有的挑战：
长期以来，信息检索技术一直是搜索引擎和质量检查任务的关键组成部分。 近来，将经典的IR技术与现代变压器相结合已经取得了令人振奋的进步，特别是用于答疑任务，其中与神经检索器一起训练模型，该神经检索器学会检索相关文档以帮助回答问题。 有关概述，请参见FAIR中的最新工作。 检索增强生成：简化智能自然语言处理模型的创建 和 解码器融合，Google Brain的 领域和Nvidia的 神经猎犬 回答问题。
尽管增强检索技术可以帮助您提高成本和效率，但我们仍然无法将模型拟合到单个GPU上以实现最大的模型。 这意味着我们需要使用模型并行性来训练它。 但是，由于我们的检索体系结构的性质，由于跨训练输入的检索上下文之间的相互依赖性，设计模型拆分是一项挑战。 此外，即使我们确定如何拆分模型，引入模型并行性也是我们在研究和开发生命周期中手动完成的一项重要工程任务。
SageMaker模型并行性库
模型并行性是在多个设备或节点（例如配备GPU的实例）之间拆分模型，并创建有效的管道以在这些设备上训练模型以最大程度地利用GPU的过程。 这 模型并行库 在SageMaker中，通过提供自动模型拆分（也称为“ 自动模型分割 以及复杂的管道运行计划。 模型拆分算法可以优化速度或内存消耗。 该库使用分区算法来平衡内存，最大程度地减少设备之间的通信并优化性能。
自动模型分割
对于我们的PyTorch用例，模型并行库在内部运行一个跟踪步骤（在第一个训练步骤中），该步骤构成模型图并确定张量和参数形状。 然后，它构造一棵树，该树由嵌套的 nn.Module 模型中的对象，以及从跟踪收集的其他数据，例如存储的数量 nn.Parameters，以及每个的运行时 nn.Module.
然后，库从根目录遍历该树并运行分区算法，该算法平衡计算负载和内存使用，并最大程度地减少实例之间的通信。 如果多个nn.Modules共享相同的nn.Parameter，则将这些模块放置在同一设备上，以避免维护同一参数的多个版本。 做出分区决定后，将分配的模块和权重加载到其设备中。
管道运行计划
SageMaker分布式模型并行库的另一个核心功能是 流水线运行，确定模型训练期间跨设备进行计算和处理数据的顺序。 流水线的基础是将微型批次分成微批，这些微批被逐一馈入训练流水，并遵循库运行时定义的运行时间表。
微批处理管道可确保充分利用所有GPU，这是我们必须自己构建的东西，但是使用模型并行库，此操作巧妙地在后台进行。 最后，我们可以使用 亚马逊FSx，这对于在训练带检索的多峰模型的过程中，鉴于要读取的文件数量非常重要，对于确保我们的读取速度非常重要。
培训架构
下图显示了我们如何设置培训体系结构。 我们的主要目标是提高培训速度并降低成本。 我们正在训练的图像和语言转换器非常复杂，内部有大量的层和权重，运行着数十亿个参数，所有这些使它们无法放入单个节点的内存中。 每个节点都承载模型的一个子集，通过该子集可以共享和编译数据流和转换。 我们设置了16 p4d.24xlarge 每个实例具有八个GPU，并使用以下架构表示形式：
在扩展模型时，一个普遍的趋势是将所有内容存储在网络的权重中。 但是，出于实际目的，我们希望扩充模型以学习如何查找相关上下文以帮助执行渲染任务。 这使我们能够在不影响图像质量的情况下降低服务成本。 我们使用了一个大型的基于变压器的NLP模型，并且如前所述，我们发现SageMaker模型并行性库的训练效率提高了38％，如下所示：
- 在张量级并行度的情况下，我们需要对所有计算进行缩减。 这需要O（log_2 n）个并行步骤。 也就是说，n台机器执行O（n）个步骤，总共执行O（n log_2 n）个操作。
- 对于管道并行性，我们需要O（1）并行步骤才能将数据向下传递到管道
- 给定16台具有1个GPU的机器，对于并行管道，我们的成本为O（2），对于深度模型并行，我们的成本为O（log_8（3））= O（XNUMX）。
- 在这种情况下，我们看到通过切换到与SageMaker模型并行性一起使用的管道并行性，网络成本降低到了1/3，并且总体培训成本降低到1/2 + 1/2 * 1 / log_2（16 ）=原始成本的0.625，因此相应地提高了效率。
通常，当需要进行分布式培训（带有缩放模型大小或培训数据的问题）时，我们可以遵循一组最佳实践来确定哪种方法最有效。
分布式培训的最佳做法
根据我们的经验，我们建议从分布式数据并行方法开始。 分布式数据并行性，例如 SageMaker分布式数据并行库 解决了模型副本的大多数网络问题，因此您应将模型放入最少数量的节点中，然后根据需要进行复制以扩展批量大小。
如果像我们在这种情况下那样在训练期间内存不足，则可能要切换到模型并行方法。 但是，在尝试模型并行训练之前，请考虑以下替代方法：
- 在配备NVIDIA Tensor Core的硬件上，使用 混合精度训练 创建加速并减少内存消耗。
- 减小批次大小（或尽可能减小图像分辨率或NLP序列长度）。
此外，我们更喜欢没有批量标准化的模型设计，如 无需归一化的高性能大规模图像识别。 如果无法避免，请确保跨设备同步批处理规范化。 使用分布式培训时，您的批次将在GPU之间划分，因此准确的批次统计信息需要在所有设备之间进行同步。 否则，归一化将增加错误，从而削弱收敛。
当您具有以下约束时，请从模型并行训练开始：
- 您的模型不适用于单个设备
- 由于模型大小，选择较大的批处理大小时会遇到限制，例如模型权重占用了大部分GPU内存，而您不得不选择较小的，次优的批处理大小
优化性能时，请执行以下操作：
- 使用流水线进行节点间通信，以最大程度地减少延迟并提高吞吐量
- 保持管线尽可能短，以最大程度地减少气泡。 微批处理的数量应进行调整，以使计算效率与气泡大小保持平衡，并且至少应为管道长度。 如果需要，您可以按照以下说明在令牌级别形成微批次 TeraPipe：令牌级管道并行性，用于训练大规模语言模型
在优化成本时，请使用SageMaker管理的竞价型实例进行培训。 在按需实例上，这可以将培训模型的成本优化高达90％。 SageMaker代表您管理Spot中断。
其他要考虑的因素：
- 在节点内，如果存在快速互连，则更加细微。 如果节点间网络容量足够大，则重新组合数据以进行更优化的计算可能会显示出好处。
- 如果激活量远大于重量张量，则分片优化器也可能会有所帮助。 请参阅 零 以获得更多细节。
下表提供了一些常见的培训扩展方案，以及如何在AWS上对其进行配置。
| EventXtra XNUMX大解决方案 | 什么时候适用？ | 解决方案 |
| 从单个GPU扩展到多个GPU | 当训练数据量或模型大小太大时 | 更改为多GPU实例（例如p3.16xlarge），该实例具有八个GPU，数据和处理过程分散在八个GPU中，并且在训练模型时产生了近乎线性的加速。 |
| 从单个实例扩展到多个实例 | 当扩展需求超出了更改实例大小的范围时 | 通过将instance_type设置为p3.16xlarge并将instance_count设置为2，可以使用SageMaker Python SDK的估算器功能扩展实例的数量。您可以在两个相同的实例上拥有3.16个GPU，而不是单个p16xlarge上的八个GPU。 考虑使用 SageMaker分布式数据并行库. |
| 选择模型并行方法进行训练 | 训练过程中遇到记忆不足的错误 | 使用 SageMaker分布式模型并行库. |
| 节点间通信的网络性能 | 对于具有多个实例的分布式训练（例如，执行AllReduce操作时集群中节点之间的通信） | 您的实例必须位于相同的区域和相同的可用区中。 当您使用SageMaker Python SDK时，将为您处理。 您的训练数据也应位于同一可用区中。 考虑使用 SageMaker分布式数据并行库. |
| 优化的GPU，网络和存储 | 满足大规模分布式培训需求 | p4d.24xlarge实例类型旨在用于快速本地存储和具有高达400吉比特的快速网络背板，我们强烈建议将其作为分布式培训中性能最高的选项。 |
结论
借助SageMaker中的模型并行库，我们可以立即获得许多好处，例如自动模型分区和有效的流水线操作。 在本文中，我们分享了ML用例的挑战，对不同培训方法的考虑以及我们如何使用Amazon SageMaker模型并行性库来加快培训速度。 最重要的是，现在只需几个小时即可采用此处介绍的模型并行性和性能改进的最佳实践。 如果这篇文章对您有帮助或启发您解决问题，我们很乐意听到！ 请分享您的评论和反馈。
参考资料
有关更多信息，请参见以下内容：
作者简介
普雷姆兰加 是总部位于乔治亚州亚特兰大的企业解决方案架构师。 他是机器学习技术领域社区的成员，并且喜欢与客户一起进行ML和AI之旅。 Prem对机器人技术充满热情，是自动驾驶汽车研究人员，还在休斯敦和其他地方建造了Alexa控制的Beer Pours。
莎拉·简·洪（Sarah Jane Hong） 是Latent Space的联合创始人兼首席科学官。 她的背景是人机交互和机器学习的交集。 她之前曾在Sonar（由Marchex收购）领导NLP研究，Sonar在对话AI领域为企业提供服务。 她还是受人尊敬的AR / VR开发人员，曾获得Oculus，Mozilla Mixed Reality和Microsoft Hololens的奖项和研究金。
达里尔·巴恩哈特（Darryl Barnhart） 是Latent Space的联合创始人兼首席技术官。 他是一位经验丰富的开发人员，在GPU加速，计算机图形，大规模数据和机器学习方面经验丰富。 其他爱好包括数学，游戏开发和信息研究。
伊恩·汤普森（Ian Thompson） 是Latent Space的创始人兼首席执行官。 伊恩（Ian）是一名工程师和研究员，受到“邻近可能”的启发-这些技术将对我们的生活产生重大影响。 目前专注于简化和扩展多模式表示学习，以帮助构建安全且富有创意的AI。 他先前曾帮助建立图形/虚拟现实（AltspaceVR，被微软收购）和教育/ NLP（HSE）公司。