如何 BRIA AI 利用 Amazon SageMaker 的分布式训练来训练用于商业用途的潜在扩

BRIA AI如何在Amazon SageMaker中使用分布式训练以训练用于商业的潜在扩散基础模型

关键要点

本文探讨了BRIA AI如何有效地在Amazon SageMaker中进行大规模模型训练，实现快速且经济地处理海量数据，提升训练效率，减少成本。以下是本文的主要内容：

模型训练特点：BRIA AI 20是一个高分辨率的文本到图像扩散模型，训练数据量达到多个PB级别。技术挑战及解决方案：文章详细说明了BRIA AI在训练过程中面临的挑战，包括操作卓越性、数据并行性、数据加载效率和训练成本。实现成果：BRIA AI通过SageMaker的支持，高效地进行模型训练，并且实现了98的GPU利用率。

在这篇文章中，我们将介绍BRIA AI如何使用Amazon SageMaker训练其高分辨率10241024的文本到图像扩散模型BRIA AI 20，利用包含PB级别的授权图片的数据集，并且做到快速和经济。Amazon SageMaker的训练任务及其分布式训练库处理了与基础设施管理相关的繁重工作。SageMaker帮助用户建立、训练和部署机器学习ML模型，提供全托管的基础设施、工具和工作流程。

BRIA AI是一个创新平台，专注于为开发者提供负责任且开放的生成性人工智能AI，仅基于与Getty Images、DepositPhotos和Alamy等合作伙伴的授权数据训练高级模型。它为大型品牌、动画和游戏工作室及营销机构提供多模态生成模型。BRIA AI强调道德来源和商业准备，其模型开放源代码、安全且优化以兼容各类技术架构。通过解决数据采购、持续模型训练和无缝技术集成的根本挑战，BRIA AI力求成为创意AI应用开发者的首选平台。

BRIA AI 20的图像生成模型现已可在AWS Marketplace上找到。

本文讨论了BRIA AI与AWS合作以应对以下关键挑战：

挑战解决方案实现大规模模型训练的操作卓越性利用SageMaker的弹性特性，保障培训集群的健康通过数据并行性减少训练时间使用SageMaker同时运行多个实例提高GPU利用率通过高效的数据加载减少模型训练成本仅为实际训练时间付费

重要的是，BRIA AI能够在保持最初使用的HuggingFace Accelerate软件栈不变的情况下使用SageMaker。因此，转向SageMaker训练无需对BRIA AI的模型实现或训练代码进行更改。随后，BRIA AI能够在SageMaker上无缝升级其软件栈及模型训练。

训练管线架构

BRIA AI的训练管线主要由两个部分构成：

数据预处理：

数据贡献者将授权的原始图像文件上传到BRIA AI的Amazon Simple Storage Service (Amazon S3)存储桶。通过使用Amazon Simple Queue Service (Amazon SQS)和AWS Lambda函数的图像预处理管线生成缺失的图像元数据，并将训练数据打包成大型webdataset文件，以便后续高效的数据流处理和数据在GPU间的分片。

模型训练：

SageMaker分配训练任务以管理训练集群并执行训练过程。使用SageMaker的快速文件模式将数据从S3流式传输到训练实例。

预训练过程中的挑战与解决方案

预训练基础模型是一项复杂的任务，面临许多挑战，如成本、性能、调度、监控，以及整个几周训练过程所需的工程专长。

挑战1：实现大规模模型训练的操作卓越性

BRIA AI利用SageMaker训练任务的弹性特性来协调训练集群并应对故障。这些特性包括集群健康检查、内置重试和任务恢复。在任务开始前，SageMaker会对GPU进行健康检查，验证英伟达集体通信库NCCL在GPU实例上的通信情况，如有必要则替换故障实例，以确保培训脚本在健康的实例集群上启动。您还可以配置SageMaker在遇到不可恢复的GPU错误时自动重试任务。重试过程中，SageMaker将用新的实例替换遇到问题的实例，重启健康实例，然后重新启动该任务，从而加快了重启及工作负载完成的速度。

为了监控训练集群，BRIA AI利用了SageMaker与Amazon CloudWatch的集成功能，包括应用日志和云监控度量CPU、GPU和网络度量。

挑战2：通过数据并行性减少训练时间

BRIA AI需在PB级别的授权图像数据集上从零开始训练稳定性扩散20模型。单个GPU的训练可能需要几个月才能完成。为了满足截止日期要求，BRIA AI使用了数据并行性，通过SageMaker的16个p4de24xlarge实例进行训练，将总训练时间减少到两周以内。分布式数据并行训练通过将数据有效地分配给多个并行训练的设备，大大加快了大型模型的训练速度，同时定期同步梯度以保持一致的共享模型。这利用了多台设备的计算能力。BRIA AI使用一组四个p4de24xlarge实例8个A100 80GB NVIDIA GPU实现了每秒有效处理2048的吞吐量。

注意 目前可以使用p548xlarge实例8个H100 80GB GPU，在使用EFA 20时实例间的网络为3200 Gbps但在此次BRIA AI的预训练中未使用。

为小规模的云外训练，BRIA AI使用了Accelerate库。随着云训练的扩展，BRIA AI得以继续使用Accelerate，通过其与SageMaker及Amazon SageMaker分布式数据并行库SMDDP的内置集成。SMDDP旨在AWS基础设施上运行，分两种方式减少通信开销：

功能说明执行AllReduce操作这是分布式训练中的关键操作，负责大量通信开销优化节点间通信完全利用AWS网络基础设施和Amazon EC2实例拓扑

请注意，SageMaker训练支持多种开源分布式训练库，例如完全分片数据并行FSDP和DeepSpeed。在其他训练工作负载中，BRIA AI使用了SageMaker中的FSDP。在这种情况下，通过使用ShardingStrategySHARDGRADOP功能，BRIA AI能够实现最佳批处理大小，加速其训练流程。

挑战3：实现高效的数据加载

BRIA AI的数据集包含数亿张需要从存储加载到GPU进行处理的图像。高效访问这大规模数据集面临许多挑战：

数据可能无法容纳在单个实例的存储中。下载多TB数据集到每个训练实例非常耗时，导致GPU无所事事。从Amazon S3复制数百万个小图像文件可能成为瓶颈，因为从S3获取对象的往返时间会累积。数据需要正确分配到各个实例之间。

BRIA AI借助SageMaker快速文件输入模式解决了这些挑战，具有以下内置功能：

流式传输：我们选择直接从Amazon S3流式传输数据到训练实例，而不是在训练开始时复制数据或使用额外的分布式文件系统，这使得训练可以同步启动，无需等待下载，且减少了将数据集适配到实例存储的需求。数据分布：快速文件模式配置为在多实例之间分片数据文件，策略为S3DataDistributionType=ShardedByS3Key。本地文件访问：快速文件模式为Amazon S3中的数据提供本地POSIX文件系统接口，使得BRIA AI的数据加载器能够像访问本地文件一样访问远程数据。将文件打包为大型容器：使用数百万个小的图像和元数据文件在从对象存储如此Amazon S3流式传输数据时会产生过重的开销。为此，BRIA AI将多个文件压缩成大型 TAR 文件容器25 GB，可高效地通过快速文件模式从S3流式传输到实例。具体而言，BRIA AI使用了WebDataset来高效加载本地数据，并采用一种策略，各个GPU通过固定种子加载随机批次，而不需要实例之间的同步。这种策略有助于消除瓶颈并保持数据加载性能快速且可预测。

有关数据加载注意事项的更多信息，请参见为Amazon SageMaker训练任务选择最佳数据源的博客文章。

挑战4：仅为实际训练时间付费

大语言模型的预训练并不是连续进行的。模型训练经常需要中断进行评估和调整。例如，模型可能会停止收敛，需要调整，或者您可能想暂停训练来测试模型、精炼数据或排查问题。这些暂停导致GPU集群长时间空闲。使用SageMaker训练任务，BRIA AI只能为其的实际活跃训练时间付费。这使BRIA AI能够以更低的成本和更高的效率训练模型。

BRIA AI的训练策略由以下三步组成，以达到最佳模型收敛：

在256256分辨率的32个GPU集群上进行初步训练渐进优化至512512分辨率的64个GPU集群在10241024分辨率的128个GPU集群上进行最终训练

每一步所需的计算量各不相同，因应用了不同的权衡，例如每个分辨率的批量大小和GPU及梯度累积的上限。这种权衡在于节省成本与模型覆盖之间的储衡。

BRIA AI通过维持一致的每秒迭代率来辅助成本计算，使得训练时间的估算更为精准。这使得可以精确确定所需的迭代次数以及每小时的训练计算成本。

BRIA AI训练GPU利用率与平均批量时间：

蘑菇加速器app官方

GPU利用率：平均超过98，这表明我们的GPU在整个训练周期内得到了最大化利用，同时数据加载器以高效率流式传输数据。每秒迭代次数：训练策略由三个步骤组成在256256分辨率上进行初步训练、渐进地优化至512512分辨率，再在10241024分辨率上进行最终训练，以达到最佳的模型收敛。每个步骤的计算量因不同的批处理大小和GPU及梯度累积的上限而有所不同，所面临的权衡也在于成本节省和模型覆盖之间。

结果样例

生成图像所使用的提示 提示1：一个时尚的男子坐在户外阶梯上，穿着绿色运动卫衣和匹配的绿色裤子，黑色鞋子，戴着太阳镜。他微笑着，头发整齐，看起来有一把短胡子。旁边放着一个棕色皮包。背景是一面砖墙和一个带白色框架的窗户。 提示2：一场富有色彩的印度婚礼。微笑的新娘穿着绣有金色图案的洋红色沙丽，手上有美丽的亨娜图案，搭配传统的金色饰品。新郎则穿着金色的舍尔瓦尼和白色的达尔希。周围点缀着花卉、蜡烛和树叶，在色彩缤纷的节日氛围中，背景满是传统的物件。 提示3：一个木制托盘上装满了各类美味的糕点。托盘里有一块撒上糖粉的可颂、一块巧克力夹心可颂、一块半吃掉的可颂、一块丹麦酥及一个旁边的小杯子装着巧克力酱，以及一碗咖啡豆，所有这些都摆在一个米色布料上。 提示4：一只熊猫在桌子上向白色杯子里倒牛奶，桌子上有咖啡豆、花和咖啡压滤器。背景是一幅黑白画和一个装饰性墙面艺术装置。

结论

在本篇文章中，我们看到了Amazon SageMaker如何帮助BRIA AI高效地训练扩散模型，而无需手动配置和管理基础设施。借助SageMaker的训练，BRIA AI能够降低成本，加快迭代速度，同时保持98的GPU利用率，最大化每一笔投入的价值。SageMaker承担了繁重的基础设施任务，使BRIA AI团队更具生产力，从而更快速地交付创新。SageMaker训练任务提供的易用性和自动化使其成为希望高效训练大规模前沿模型的团队的理想选择。

想了解更多关于如何使用SageMaker高效且经济地训练大型AI模型的信息，请访问Amazon SageMaker页面。您还可以与您的AWS帐户团队联系，了解如何释放全规模AI项目的潜力。

关于作者

Bar Fingerman，BRIA AI人工智能/机器学习工程负责人。

Doron Bleiberg，高级初创解决方案架构师。

Gili Nachum，首席生成性AI/机器学习专属解决方案架构师。

Erez Zarum，初创解决方案架构师。