開(kāi)源的 Sora 復(fù)現(xiàn)方案,成本降低近一半!

共 2362字，需瀏覽 5分鐘

2024-04-11 07:00

近日，開(kāi)發(fā) ChatGPT 的 OpenAI 公司又放出王炸 Sora，一個(gè)可以根據(jù)文本生成視頻的 AI 模型。

上圖就是 OpenAI 公布的 Sora 生成的視頻片段，可以毫不夸張地說(shuō) Sora 直接將視頻生成技術(shù)推向了新的高度，這也標(biāo)志著人工智能視頻生成技術(shù)邁入了新的時(shí)代。此項(xiàng)技術(shù)，可以廣泛應(yīng)用于電影、動(dòng)畫(huà)、游戲、廣告等領(lǐng)域，為內(nèi)容創(chuàng)作者提供更加便捷、高效的創(chuàng)作工具。

雖然 Sora 沒(méi)有開(kāi)源，但我在 GitHub 上發(fā)現(xiàn)了 Colossal-AI 剛開(kāi)源的完整 Sora 復(fù)現(xiàn)架構(gòu)方案 Open-Sora，它提供了完整的 Sora 復(fù)現(xiàn)架構(gòu)解決方案，包括從數(shù)據(jù)處理到訓(xùn)練、部署的全流程。Open-Sora 不僅可以降低 46% 復(fù)現(xiàn)成本，而且還將模型訓(xùn)練輸入序列長(zhǎng)度擴(kuò)充至 819K patches（模型處理視頻的最小單元）。

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

既然 Sora 沒(méi)有開(kāi)源，那這個(gè)復(fù)現(xiàn)方案從何而來(lái)呢？接下來(lái)，就讓我們一起來(lái)看看已公布的 Sora 技術(shù)原理以及 Open-Sora 到底有沒(méi)有真東西！

Sora 算法復(fù)現(xiàn)方案

與 Sora 視頻一起發(fā)布的，還有一份 OpenAI 關(guān)于 Sora 的技術(shù)報(bào)告，在這份 Sora 的技術(shù)報(bào)告中，介紹了 Sora 是使用了一個(gè)視頻壓縮網(wǎng)絡(luò)，將不同尺寸的視頻壓縮成一個(gè)隱空間（latent space）的時(shí)空塊序列（temporal patch），然后使用了 Diffusion Transformer 進(jìn)行去噪，最后進(jìn)行解碼生成視頻。

根據(jù)上面描述的技術(shù)原理，Open-Sora 將 Sora 可能使用的訓(xùn)練流程歸納為下圖。

綜上所述，Open-Sora 作為 Sora 技術(shù)的復(fù)現(xiàn)框架，用 Python 實(shí)現(xiàn)了將原視頻轉(zhuǎn)化成通用型大模型能處理的最小單元 patches 的模塊（patches 類(lèi)似文本的 token），然后在去噪步驟提供了 3 種常見(jiàn)的多模態(tài)模型結(jié)構(gòu)。

目前 Open-Sora 提供的功能，如下：

完整的 Sora 復(fù)現(xiàn)架構(gòu)：包含從數(shù)據(jù)處理到訓(xùn)練推理全流程。
動(dòng)態(tài)分辨率：訓(xùn)練時(shí)可直接訓(xùn)練任意分辨率的視頻，無(wú)需進(jìn)行縮放。
多種模型結(jié)構(gòu)：由于 Sora 實(shí)際模型結(jié)構(gòu)未知，我們實(shí)現(xiàn)了 adaLN-zero、cross attention、in-context conditioning（token concat）等 3 種常見(jiàn)的多模態(tài)模型結(jié)構(gòu)。
多種視頻壓縮方法：用戶(hù)可自行選擇使用原始視頻、VQVAE（視頻原生的模型）、SD-VAE（圖像原生的模型）進(jìn)行訓(xùn)練。
多種并行訓(xùn)練優(yōu)化：支持結(jié)合 Colossal-AI 的 AI 大模型系統(tǒng)優(yōu)化能力，及 Ulysses 和 FastSeq 的混合序列并行。

眾所周知，模型訓(xùn)練需要的資源是恐怖的，哪怕是 1% 的優(yōu)化都可以節(jié)約極大的訓(xùn)練成本，而 Open-Sora 直接節(jié)省了近一半 46% 的復(fù)現(xiàn)成本，而且在模型訓(xùn)練輸入方面更是將序列擴(kuò)充至近百萬(wàn)。到底 Open-Sora 是如何做的性能優(yōu)化，從而實(shí)現(xiàn)類(lèi) Sora 視頻生成模型的低成本開(kāi)發(fā)流程呢？跟著小編我們一探究竟。

性能優(yōu)化

不同于 LLM 的大模型、大激活，Sora 類(lèi)訓(xùn)練任務(wù)的特點(diǎn)是模型本體不大（如在 10B 以下），但是由于視頻復(fù)雜性帶來(lái)的序列長(zhǎng)度特別長(zhǎng)。

在此情況下，PyTorch 數(shù)據(jù)并行已無(wú)法運(yùn)行，而傳統(tǒng)的模型并行、零冗余數(shù)據(jù)并行帶來(lái)的收益有限。因此，在支持 AMP（FP16/BF16）、Flash Attention、Gradient checkpointing、ZeRO-DP 等場(chǎng)景優(yōu)化策略的基礎(chǔ)上，Open-Sora 進(jìn)一步引入兩種不同的序列并行方法實(shí)現(xiàn)，可以 ZeRO 一起使用實(shí)現(xiàn)混合并行：

通用性較強(qiáng)的 Ulysses，對(duì)小規(guī)?；蜷L(zhǎng)序列表現(xiàn)可能更好。
FastSeq 能將 qkv projection 的計(jì)算和 all-gather 通信重疊，只需多占用一點(diǎn)內(nèi)存就可更進(jìn)一步提升訓(xùn)練效率。

這兩種序列并行方案，都可以輕松與 ZeRO2 共同使用來(lái)實(shí)現(xiàn)混合并行。

以在單臺(tái) H800 SXM 8*80GB GPU 上使用 DiT-XL/2 模型的性能測(cè)試為例。

在 600K 的序列長(zhǎng)度時(shí)，Open-Sora 的方案比基線(xiàn)方案有 40% 以上的性能提升和成本降低。

在保證更快訓(xùn)練速度的情況下，Open-Sora 還能訓(xùn)練 30% 更長(zhǎng)的序列，達(dá)到 819K+。

最后

在這個(gè)視頻流行的時(shí)代，AI 生成視頻技術(shù)成為了當(dāng)下備受矚目的熱門(mén)技術(shù)之一，Sora 的出現(xiàn)更是引爆了這個(gè)領(lǐng)域。

今天介紹的 Open-Sora 才剛剛開(kāi)源，并未提供訓(xùn)練好的類(lèi) Sora 模型，還不能直接拿來(lái)生成視頻。但它提供了一套經(jīng)過(guò)優(yōu)化后的類(lèi) Sora 視頻生成模型低成本開(kāi)發(fā)框架，為視頻生成提供方便易用、成本低廉、質(zhì)量可靠的開(kāi)源解決方案。更多的技術(shù)細(xì)節(jié)，篇幅問(wèn)題這里就不展開(kāi)了，如果你對(duì) Open-Sora 感興趣，可深入了解下它的源碼（Python）或者參與貢獻(xiàn)代碼，復(fù)現(xiàn) Sora 趕超 Sora！