開源的 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

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章