可扩展的训练平台:使用 Amazon SageMaker HyperPod 进行创新视频生成案例
- 20
使用 Amazon SageMaker HyperPod 的可扩展培训平台:视频生成案例研究
关键要点
视频生成是人工智能研究中的新兴领域,Luma AI 推出的 Dream Machine 拉开了新篇章。Amazon SageMaker HyperPod 提供了一种专门构建的基础设施,旨在解决大规模培训所面临的挑战。通过 SageMaker HyperPod,用户可以快速迭代,提升视频生成模型的效率和性能。在人工智能研究领域,视频生成已经成为最新的前沿技术,紧随文本到图像模型的成功。Luma AI 最近推出的 Dream Machine 在这个领域代表了一个重要的进步。该文本到视频 API 可以快速从文本和图像生成高质量、真实感的的视频。Dream Machine 由 Amazon SageMaker HyperPod 训练,擅长创建一致的角色、流畅的运动和动态摄像机移动。
为了加速该领域的迭代和创新,充足的计算资源和可扩展的平台至关重要。在迭代研究和开发阶段,数据科学家和研究人员需要执行多个实验,以不同版本的算法进行测试,并扩展到更大的模型。当模型总占用内存超出单个 GPU 的内存限制时,模型并行训练变得必要。然而,构建大型分布式训练集群是一个复杂且耗时的过程,需要深入的专业知识。此外,随着集群规模的扩大例如,超过32个节点,它们还需要内置的恢复机制,如自动故障节点检测和替换,以提高集群的良品率并维持高效的操作。这些挑战突显了在支持先进的 AI 研发中,强大基础设施和管理系统的重要性。
在 reInvent 2023 上推出的 Amazon SageMaker HyperPod,是一种专门设计的基础设施,旨在应对大规模培训的挑战。它消除了构建和优化机器学习基础设施所需的繁重工作,使用户能够专注于算法创新。SageMaker HyperPod 提供了高度可定制的用户界面,使用 Slurm,允许用户选择和安装所需的框架或工具。集群可以根据用户的选择进行配置,并能够跨工作负载保留。凭借这一能力,客户将 SageMaker HyperPod 作为更具弹性和更高性能模型培训的创新平台,使他们能够更快构建最先进的模型。
在本文中,我们将分享一个使用 SageMaker HyperPod 支持视频生成研究团队创新的机器学习基础设施架构。我们将讨论 SageMaker HyperPod 所解决的问题及其优点,提供逐步设置指南,并展示如何在集群上运行视频生成算法。
在 Amazon SageMaker HyperPod 上培训视频生成算法:背景与架构
视频生成是一个令人兴奋且快速演变的领域,近年来取得了显著进展。尽管生成建模在图像生成领域取得了巨大进展,视频生成仍面临多项挑战,需要进一步改进。
扩散模型家族的算法架构复杂性
最近,扩散模型在生成高质量图像方面取得了显著进展,促使研究人员探索其在视频生成中的潜力。利用扩散模型的架构及预训练生成能力,科学家们旨在创造视觉上令人印象深刻的视频。该过程将图像生成技术扩展到时间域。模型从噪声帧开始,迭代地进行精炼,移除随机元素,同时根据文本或图像提示添加有意义的细节。这一方法逐步将抽象模式转变为连贯的视频序列,有效地将扩散模型在静态图像创作中的成功转化为动态视频合成。
然而,使用扩散模型进行视频生成的计算需求大幅增加,原因如下:
时间维度:与图像生成不同,视频生成需要同时处理多个帧,为原有的 2D UNet 增加了时间维度,显著增加了需要并行处理的数据量。迭代去噪过程:扩散过程需要对每个帧进行多次迭代去噪。扩展到视频时,这个迭代过程必须应用于多个帧,导致计算负担加倍。参数数量增加:为了处理视频数据的额外复杂性,模型通常需要更多参数,从而导致内存占用和计算需求增加。更高的分辨率和更长的视频序列:视频生成通常旨在输出高分辨率和更长序列,相较于单张图像生成,计算要求进一步上升。由于这些因素,扩散模型在视频生成中的运行效率较低,计算强度显著高于图像生成。这一增加的计算需求强调了需要先进的硬件解决方案和优化模型架构,以使视频生成变得更为实用。
应对增加的计算需求
视频生成质量的提高需要显著增加模型及训练数据的规模。研究人员已得出结论,扩展基础模型的规模会显著提升视频生成性能。然而,这一增长带来了计算能力和内存资源方面的挑战。培训更大的模型需要更多的计算能力和内存空间,这可能限制这些模型的可用性和实用性。随着模型规模的增加,计算需求呈指数增长,难以在单个 GPU 或单节点多 GPU 环境中进行训练。此外,存储和处理所需的大型训练数据集在基础设施和成本方面也面临重大挑战。高质量的视频数据集往往庞大,需要巨大的存储容量和高效的数据管理系统。转移和处理这些数据集可能耗时繁琐,进一步增加了整体计算负担。
保持时间一致性和连续性
随着生成视频的长度增加,保持时间一致性和连续性变得日益具有挑战性。时间一致性指的是视觉元素例如物体、角色和场景在后续帧中的连续性。外观、运动或照明上的不一致会导致视觉伪影,打断整体观看体验。为了解决这一挑战,研究人员探索了使用多帧输入的技术,这为模型提供了来自多个连续帧的信息,以更好地理解和建模时间上的关系与依赖。这些技术在保持高分辨率细节的同时,模拟连续和平滑的时间运动过程。然而,这些技术需要更为复杂的建模技术和增加的计算资源。
算法概述
在接下来的部分中,我们将演示如何在 Amazon SageMaker HyperPod 上运行 Animate Anyone Consistent and Controllable ImagetoVideo Synthesis for Character Animation 算法进行视频生成。Animate Anyone 是一种将角色图像转换为根据所需姿势序列控制的动画视频的方法。架构的关键组成部分包括:
ReferenceNet:一种对称的 UNet 结构,捕获参考图像的空间细节,并将其集成到去噪 UNet 中,利用空间注意力来保持外观一致性。Pose guider:一个轻量级模块,有效地在去噪过程中集成姿势控制信号,以确保姿势的可控性。Temporal layer:添加到去噪 UNet 中,以建模多个帧之间的关系,保持高分辨率细节,确保角色运动的时间稳定性和连续性。模型架构如下面的图像所示,来自其 原始研究论文。该方法在一组视频剪辑数据集上进行训练,在时尚视频和人舞合成基准上取得最先进的结果,展示了其能够在保持外观一致性和时间稳定性的同时,对任意角色进行动画处理。AnimateAnyone 的实现可在 这个仓库 中找到。
为了应对视频生成训练过程中所需的大规模培训基础设施所带来的挑战,我们可以利用 Amazon SageMaker HyperPod 的强大功能。许多客户已经采用 SageMaker HyperPod 进行大规模培训,如 Luma 的 Dream Machine 推出 和 Stability AI 在图像或视频生成方面的工作,我们相信 SageMaker HyperPod 的能力也能惠及更轻量的机器学习工作负载,包括完整的微调。
Amazon SageMaker HyperPod 概念与优势
SageMaker HyperPod 提供了一组全面的功能,显著提高了机器学习工作流程的效率和效果。从支持分布式培训的专用基础设施到可定制的环境,以及与 Slurm 等工具的无缝集成,SageMaker HyperPod 使机器学习从业者能够专注于核心任务,同时利用分布式计算的强大能力。使用 SageMaker HyperPod,您可以加快机器学习项目的进展,处理更大的数据集和模型,并推动组织内部的创新。SageMaker HyperPod 在可扩展培训架构中的多个关键特性和优势如下表所示。
加速器下载官方特性优势专用基础设施简化了集群的设置和管理,使用户能够轻松配置所需的实例类型和数量共享文件系统支持与 Amazon FSx for Lustre 的集成,提供高效的数据存储和获取可定制环境支持使用生命周期脚本自定义集群环境,允许用户按需安装框架和工具与 Slurm 集成无缝集成 Slurm,允许高效调度工作,提高资源使用效率增强生产力通过 SSH 连接使用 Visual Studio Code 连接到集群,提高开发效率设置 SageMaker HyperPod 并运行视频生成算法
在本指南中,我们将使用 AnimateAnyone 算法作为视频生成的示例。AnimateAnyone 是一种从输入图像或视频中生成高质量视频的最先进算法。我们的逐步指导代码可在 GitHub 上获取。
设置集群
要创建 SageMaker HyperPod 基础设施,请按照 Amazon SageMaker HyperPod 研讨会工作室 的详细步骤进行配置。
您需要准备两件事:一个由 HyperPod 设置 Slurm 所需的 provisioningparametersjson 文件,以及作为创建 HyperPod 集群的配置文件的 clusterconfigjson 文件。在这些配置文件中,您需要指定控制组和工作组的 InstanceGroupName、InstanceType 和 InstanceCount,以及附加到该组的执行角色。
一个实用的设置是在 Amazon FSx 和 Amazon S3 之间建立 双向同步 。这可以通过 Amazon S3 与 Amazon FSx for Lustre 的集成来实现,帮助在文件系统和 Amazon S3 之间建立完整的双向同步。此外,它还能将文件系统与多个 S3 存储桶或前缀同步。
此外,如果您更喜欢使用如 VSCode 这样的本地 IDE,您可以在 IDE 中设置 SSH 连接到控制节点。这样,工作节点可以用于在 Conda 环境内运行脚本和 Jupyter Notebook 服务器。
运行 AnimateAnyone 算法
当集群投入使用时,您可以通过 SSH 连接到控制节点,然后进入工作节点,访问可用的 GPU 计算资源。您可以按照 SSH 访问计算资源 的指南进行操作。建议直接在工作节点上安装所需库。
要创建 Conda 环境,请按照 Miniconda 的 快速命令行安装 的说明进行操作。然后您可以使用该 Conda 环境安装所有必需库。
bashsource /miniconda3/bin/activateconda create n videogen pip install r requirementstxt
要运行 AnimateAnyone,克隆 GitHub 仓库并按照说明进行操作。
要训练 AnimateAnyone,首先启动第一阶段以训练去噪 UNet 和 ReferenceNet,使模型能够在给定参考图像和目标姿势的条件下生成高质量的动画图像。去噪 UNet 和 ReferenceNet 是基于来自 Stable Diffusion 的预训练权重进行初始化的。
bashaccelerate launch trainstage1py config configs/train/stage1yaml
第二阶段的目标是训练时间层,以捕获视频帧之间的时间依赖性。
bashaccelerate launch trainstage2py config configs/train/stage2yaml
一旦训练脚本按预期执行,可以使用 Slurm 调度作业在单个节点上运行。我们提供了一个批处理文件来模拟单节点训练工作。可以是单 GPU 或具有多个 GPU 的单节点。如果您想了解更多,文档提供了关于在 SageMaker HyperPod 集群上运行工作的详细说明。
bashsbatch submitanimateanyonealgosh
!/bin/bash
SBATCH jobname=videogen
SBATCH N 1
SBATCH exclusive
SBATCH o videogenstage1out
export OMPNUMTHREADS=1
激活 conda 环境
source /miniconda3/bin/activateconda activate videogensrun accelerate launch trainstage1py config configs/train/stage1yaml

通过以下代码片段检查作业状态。
bashsqueueJOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)10 dev videoge ubuntu R 016 1 ip10193196
通过使用较小的批量大小并设置 use8bitadam=True,您可以在单个 GPU 上实现高效训练。当使用单 GPU 时,建议使用多 GPU 集群来进行多个实验。
以下代码块是并行运行四个作业以测试不同超参数的一个示例。我们在此提供批处理文件。
bashsbatch submithyperparametertestingsh
squeueJOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)40 dev videoge ubuntu R 008 1 ip101175641 dev videoge ubuntu R 008 1 ip101334942 dev videoge ubuntu R 008 1 ip1013715243 dev videoge ubuntu R 008 1 ip1018368
然后可以比较实验,并选择最佳配置。在我们的场景中,如下图所示,我们使用不同的数据集和视频预处理策略验证第一阶段的训练。然后,我们快速得出有关视频质量对第一阶段训练结果的影响的结论。对于实验跟踪,除了在控制节点上安装 MLflow 以监控训练进展,您还可以利用 Amazon SageMaker 上完全托管的 MLflow 能力。这样,数据科学家可以轻松地在 SageMaker 上使用 MLflow 进行模型训练、注册和部署。
扩展到多节点 GPU 设置
随着模型尺寸的增长,单个 GPU 内存很快就会成为瓶颈。大模型容易因为普通数据并行而耗尽内存,实施模型并行可能会面临挑战。DeepSpeed 解决了这些问题,加快了模型的开发和训练。
ZeRODeepSpeed 是一个深度学习优化库,其目标是简化分布式训练,使其高效有效。DeepSpeed 的 ZeRO 通过在数据并行进程之间分区三个模型状态优化器状态、梯度和参数,消除了内存冗余,而不是复制它们。这种方法显著提高了内存效率,而不降低计算粒度和通信效率。
ZeRO 提供了三个阶段的优化:
ZeRO Stage 1:在进程之间分区优化器状态,每个进程仅更新其分区。ZeRO Stage 2:进一步分区梯度,每个进程仅保留与其优化器状态部分相关的梯度。ZeRO Stage 3:在进程之间分区模型参数,在前向和反向传递过程中自动收集和分区它们。每个阶段提供的内存效率逐步提高,但通信开销也相应增加。这些技术使得训练极大的模型成为可能,否则将是无法实现的。这在有限 GPU 内存或训练非常大模型时尤其有用。
AccelerateAccelerate 是一个库,使得在任何分布式配置下运行相同的 PyTorch 代码变得简单,