Colossal-AI开源了完整的Sora复现架构方案Open-Sora,声称可降低46%复现成本,并将模型训练输入序列长度扩充至819K patches。
Sora 算法复现方案
在Sora的技术报告中,Sora使用了一个视频压缩网络将各种尺寸的视频压缩成一个隐空间(latent space)的时空块序列(a sequence of patial temporal patch),然后使用了Diffusion Transformer进行去噪,最后进行解码生成视频。
Open-Sora将Sora可能使用的训练pipeline归纳为下图。
目前Open-Sora已涵盖:
-
提供完整的Sora复现架构方案,包含从数据处理到训练推理全流程。
-
支持动态分辨率,训练时可直接训练任意分辨率的视频,无需进行缩放。
-
支持多种模型结构。由于Sora实际模型结构未知,我们实现了adaLN-zero、cross attention、in-context conditioning(token concat)等三种常见的多模态模型结构。
-
支持多种视频压缩方法。用户可自行选择使用原始视频、VQVAE(视频原生的模型)、SD-VAE(图像原生的模型)进行训练。
-
支持多种并行训练优化。包括结合Colossal-AI的AI大模型系统优化能力,及Ulysses和FastSeq的混合序列并行。
性能
以在单台H800 SXM 8*80GB GPU上使用DiT-XL/2模型的性能测试为例。在600K的序列长度时,Open-Sora的方案比基线方案有40%以上的性能提升和成本降低。
Open-Sora 开源地址:https://github.com/hpcaitech/Open-Sora