开源的 Sora 复现方案,成本降低近一半!

近日,开发 ChatGPT 的 OpenAI 公司又放出王炸 Sora,一个可以根据文本生成视频的 AI 模型。

上图就是 OpenAI 公布的 Sora 生成的视频片段,可以毫不夸张地说 Sora 直接将视频生成技术推向了新的高度,这也标志着人工智能视频生成技术迈入了新的时代。此项技术,可以广泛应用于电影、动画、游戏、广告等领域,为内容创作者提供更加便捷、高效的创作工具。

虽然 Sora 没有开源,但我在 GitHub 上发现了 Colossal-AI 刚开源的完整 Sora 复现架构方案 Open-Sora,它提供了完整的 Sora 复现架构解决方案,包括从数据处理到训练、部署的全流程。Open-Sora 不仅可以降低 46% 复现成本,而且还将模型训练输入序列长度扩充至 819K patches(模型处理视频的最小单元)。

GitHub 地址:https://github.com/hpcaitech/Open-Sora

既然 Sora 没有开源,那这个复现方案从何而来呢?接下来,就让我们一起来看看已公布的 Sora 技术原理以及 Open-Sora 到底有没有真东西!

Sora 算法复现方案

与 Sora 视频一起发布的,还有一份 OpenAI 关于 Sora 的技术报告,在这份 Sora 的技术报告中,介绍了 Sora 是使用了一个视频压缩网络,将不同尺寸的视频压缩成一个隐空间(latent space)的时空块序列(temporal patch),然后使用了 Diffusion Transformer 进行去噪,最后进行解码生成视频。

根据上面描述的技术原理,Open-Sora 将 Sora 可能使用的训练流程归纳为下图。

综上所述,Open-Sora 作为 Sora 技术的复现框架,用 Python 实现了将原视频转化成通用型大模型能处理的最小单元 patches 的模块(patches 类似文本的 token),然后在去噪步骤提供了 3 种常见的多模态模型结构。

目前 Open-Sora 提供的功能,如下:

  • 完整的 Sora 复现架构:包含从数据处理到训练推理全流程。
  • 动态分辨率:训练时可直接训练任意分辨率的视频,无需进行缩放。
  • 多种模型结构:由于 Sora 实际模型结构未知,我们实现了 adaLN-zero、cross attention、in-context conditioning(token concat)等 3 种常见的多模态模型结构。
  • 多种视频压缩方法:用户可自行选择使用原始视频、VQVAE(视频原生的模型)、SD-VAE(图像原生的模型)进行训练。
  • 多种并行训练优化:支持结合 Colossal-AI 的 AI 大模型系统优化能力,及 Ulysses 和 FastSeq 的混合序列并行。

众所周知,模型训练需要的资源是恐怖的,哪怕是 1% 的优化都可以节约极大的训练成本,而 Open-Sora 直接节省了近一半 46% 的复现成本,而且在模型训练输入方面更是将序列扩充至近百万。到底 Open-Sora 是如何做的性能优化,从而实现类 Sora 视频生成模型的低成本开发流程呢?跟着小编我们一探究竟。

性能优化

不同于 LLM 的大模型、大激活,Sora 类训练任务的特点是模型本体不大(如在 10B 以下),但是由于视频复杂性带来的序列长度特别长。

在此情况下,PyTorch 数据并行已无法运行,而传统的模型并行、零冗余数据并行带来的收益有限。因此,在支持 AMP(FP16/BF16)、Flash Attention、Gradient checkpointing、ZeRO-DP 等场景优化策略的基础上,Open-Sora 进一步引入两种不同的序列并行方法实现,可以 ZeRO 一起使用实现混合并行:

  1. 通用性较强的 Ulysses,对小规模或长序列表现可能更好。
  2. FastSeq 能将 qkv projection 的计算和 all-gather 通信重叠,只需多占用一点内存就可更进一步提升训练效率。

这两种序列并行方案,都可以轻松与 ZeRO2 共同使用来实现混合并行。

以在单台 H800 SXM 8*80GB GPU 上使用 DiT-XL/2 模型的性能测试为例。

在 600K 的序列长度时,Open-Sora 的方案比基线方案有 40% 以上的性能提升和成本降低

在保证更快训练速度的情况下,Open-Sora 还能训练 30% 更长的序列,达到 819K+

最后

在这个视频流行的时代,AI 生成视频技术成为了当下备受瞩目的热门技术之一,Sora 的出现更是引爆了这个领域。

今天介绍的 Open-Sora 才刚刚开源,并未提供训练好的类 Sora 模型,还不能直接拿来生成视频。但它提供了一套经过优化后的类 Sora 视频生成模型低成本开发框架,为视频生成提供方便易用、成本低廉、质量可靠的开源解决方案。更多的技术细节,篇幅问题这里就不展开了,如果你对 Open-Sora 感兴趣,可深入了解下它的源码(Python)或者参与贡献代码,复现 Sora 赶超 Sora!

GitHub 地址:https://github.com/hpcaitech/Open-Sora