Open-Sora全面開源模型參數和所有訓練細節

不久前OpenAI Sora以其驚人的視頻生成效果迅速走紅，在一眾文生視頻模型中突出重圍，成為全球矚目的焦點。繼2周前推出成本直降46%的Sora訓練推理復現流程后，Colossal-AI團隊全面開源全球首個類Sora架構視頻生成模型 「Open-Sora 1.0」，涵蓋了整個訓練流程，包括數據處理、所有訓練細節和模型權重，攜手全球AI熱愛者共同推進視頻創作的新紀元。

先睹為快，我們先看一段由Colossal-AI團隊發布的「Open-Sora 1.0」模型生成的都市繁華掠影視頻。

這僅僅是Sora復現技術冰山的一角，關于以上文生視頻的模型架構、訓練好的模型權重、復現的所有訓練細節、數據預處理過程、demo展示 和 詳細的上手教程，Colossal-AI團隊已經全面免費開源在GitHub，同時筆者第一時間聯系了該團隊，了解到他們將不斷更新Open-Sora的相關解決方案和最新動態，感興趣的朋友可以持續關注Open-Sora的開源社區。

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

全面解讀Sora復現方案

接下來，我們將深入解讀Sora復現方案的多個關鍵維度，包括模型架構設計、訓練復現方案、數據預處理、模型生成效果展示以及高效訓練優化策略。

模型架構設計

模型采用了目前火熱的Diffusion Transformer (DiT) [1] 架構。作者團隊以同樣使用DiT架構的高質量開源文生圖模型PixArt-α [2] 為基座，在此基礎上引入時間注意力層，將其擴展到了視頻數據上。具體來說，整個架構包括一個預訓練好的VAE，一個文本編碼器，和一個利用空間-時間注意力機制的 STDiT (Spatial Temporal Diffusion Transformer)模型。其中，STDiT 每層的結構如下圖所示。它采用串行的方式在二維的空間注意力模塊上疊加一維的時間注意力模塊，用于建模時序關系。在時間注意力模塊之后，交叉注意力模塊用于對齊文本的語意。與全注意力機制相比，這樣的結構大大降低了訓練和推理開銷。與同樣使用空間-時間注意力機制的 Latte [3] 模型相比，STDiT 可以更好的利用已經預訓練好的圖像 DiT 的權重，從而在視頻數據上繼續訓練。

STDiT結構示意圖

整個模型的訓練和推理流程如下。據了解，在訓練階段首先采用預訓練好的Variational Autoencoder (VAE)的編碼器將視頻數據進行壓縮，然后在壓縮之后的潛在空間中與文本嵌入(text embedding)一起訓練STDiT擴散模型。在推理階段，從VAE的潛在空間中隨機采樣出一個高斯噪聲，與提示詞嵌入(prompt embedding)一起輸入到STDiT中，得到去噪之后的特征，最后輸入到VAE的解碼器，解碼得到視頻。

模型的訓練流程

訓練復現方案

我們向該團隊了解到，Open-Sora 的復現方案參考了Stable Video Diffusion (SVD)[3]工作，共包括三個階段，分別是：

1) 大規模圖像預訓練；

2) 大規模視頻預訓練；

3) 高質量視頻數據微調。

每個階段都會基于前一個階段的權重繼續訓練。相比于從零開始單階段訓練，多階段訓練通過逐步擴展數據，更高效地達成高質量視頻生成的目標。

訓練方案三階段

• 第一階段：大規模圖像預訓練

第一階段通過大規模圖像預訓練，借助成熟的文生圖模型，有效降低視頻預訓練成本。

作者團隊向我們透露，通過互聯網上豐富的大規模圖像數據和先進的文生圖技術，我們可以訓練一個高質量的文生圖模型，該模型將作為下一階段視頻預訓練的初始化權重。同時，由于目前沒有高質量的時空 VAE，他們采用了Stable Diffusion [5] 模型預訓練好的圖像VAE。該策略不僅保障了初始模型的優越性能，還顯著降低了視頻預訓練的整體成本。

• 第二階段：大規模視頻預訓練

第二階段執行大規模視頻預訓練，增加模型泛化能力，有效掌握視頻的時間序列關聯。

我們了解到，這個階段需要使用大量視頻數據訓練，保證視頻題材的多樣性，從而增加模型的泛化能力。第二階段的模型在第一階段文生圖模型的基礎上加入了時序注意力模塊，用于學習視頻中的時序關系。其余模塊與第一階段保持一致，并加載第一階段權重作為初始化，同時初始化時序注意力模塊輸出為零，以達到更高效更快速的收斂。Colossal-AI團隊使用了PixArt-alpha[2]的開源權重作為第二階段STDiT模型的初始化，以及采用了T5 [6]模型作為文本編碼器。同時他們采用了256x256的小分辨率進行預訓練，進一步增加了收斂速度，降低訓練成本。

• 第三階段：高質量視頻數據微調

第三階段對高質量視頻數據進行微調，顯著提升視頻生成的質量。

作者團隊提及第三階段用到的視頻數據規模比第二階段要少一個量級，但是視頻的時長、分辨率和質量都更高。通過這種方式進行微調，他們實現了視頻生成從短到長、從低分辨率到高分辨率、從低保真度到高保真度的高效擴展。

作者團隊表示，在Open-Sora的復現流程中，他們使用了64塊H800進行訓練。第二階段的訓練量一共是 2808 GPU hours，約合7000美元，第三階段的訓練量是1920 GPU hours，大約4500美元。經過初步估算，整個訓練方案成功把Open-Sora復現流程控制在了1萬美元左右。

數據預處理

為了進一步降低Sora復現的門檻和復雜度，Colossal-AI團隊在代碼倉庫中還提供了便捷的視頻數據預處理腳本，讓大家可以輕松啟動Sora復現預訓練，包括公開視頻數據集下載，長視頻根據鏡頭連續性分割為短視頻片段，使用開源大語言模型LLaVA[7]生成精細的提示詞。作者團隊提到他們提供的批量視頻標題生成代碼可以用兩卡 3 秒標注一個視頻，并且質量接近于 GPT-4V。最終得到的視頻/文本對可直接用于訓練。借助他們在GitHub上提供的開源代碼，我們可以輕松地在自己的數據集上快速生成訓練所需的視頻/文本對，顯著降低了啟動Sora復現項目的技術門檻和前期準備。

基于數據預處理腳本自動生成的視頻/文本

模型生成效果展示

下面我們來看一下Open-Sora實際視頻生成效果。比如讓Open-Sora生成一段在懸崖海岸邊，海水拍打著巖石的航拍畫面。

再讓Open-Sora去捕捉山川瀑布從懸崖上澎湃而下，最終匯入湖泊的宏偉鳥瞰畫面。

除了上天還能入海，簡單輸入prompt，讓Open-Sora生成了一段水中世界的鏡頭，鏡頭中一只海龜在珊瑚礁間悠然游弋。

Open-Sora還能通過延時攝影的手法，向我們展現了繁星閃爍的銀河。

如果你還有更多視頻生成的有趣想法，可以訪問Open-Sora開源社區獲取模型權重進行免費的體驗。鏈接：https://github.com/hpcaitech/Open-Sora

值得注意的是，作者團隊在 Github 上提到目前版本僅使用了 400K 的訓練數據，模型的生成質量和遵循文本的能力都有待提升。例如在上面的烏龜視頻中，生成的烏龜多了一只腳。Open-Sora 1.0 也并不擅長生成人像和復雜畫面。作者團隊在Github上列舉了一系列待做規劃，旨在不斷解決現有缺陷，提升生成質量。

高效訓練加持

除了大幅降低Sora復現的技術門檻，提升視頻生成在時長、分辨率、內容等多個維度的質量，作者團隊還提供了Colossal-AI加速系統進行Sora復現的高效訓練加持。通過算子優化和混合并行等高效訓練策略，在處理64幀、512x512分辨率視頻的訓練中，實現了1.55倍的加速效果。同時，得益于Colossal-AI的異構內存管理系統，在單臺服務器上（8*H800）可以無阻礙地進行1分鐘的1080p高清視頻訓練任務。 

此外，在作者團隊的報告中，我們也發現STDiT模型架構在訓練時也展現出卓越的高效性。和采用全注意力機制的DiT相比，隨著幀數的增加，STDiT實現了高達5倍的加速效果，這在處理長視頻序列等現實任務中尤為關鍵。

歡迎持續關注Open-Sora開源項目：https://github.com/hpcaitech/Open-Sora

作者團隊提及，他們將會繼續維護和優化Open-Sora項目，預計將使用更多的視頻訓練數據，以生成更高質量、更長時長的視頻內容，并支持多分辨率特性，切實推進AI技術在電影、游戲、廣告等領域的落地。

Reference

[1] https://arxiv.org/abs/2212.09748 Scalable Diffusion Models with Transformers

[2] https://arxiv.org/abs/2310.00426 PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis

[3] https://arxiv.org/abs/2311.15127 Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets

[4] https://arxiv.org/abs/2401.03048 Latte: Latent Diffusion Transformer for Video Generation

[5] https://huggingface.co/stabilityai/sd-vae-ft-mse-original

[6] https://github.com/google-research/text-to-text-transfer-transformer

[7] https://github.com/haotian-liu/LLaVA

[8] https://hpc-ai.com/blog/open-sora-v1.0

雷峰網(公眾號：雷峰網)