Sora 在 2024 年初的驚人表現成為了新的標桿，激勵所有研究文生影片的人士爭相追趕。每位研究者都懷著復現 Sora 成果的渴望，爭分奪秒地努力著。
根據OpenAI 揭露的技術報告，Sora 的一個重要創新點是將視覺數據轉換為patch 的統一表示形式，並透過Transformer 和擴散模型相結合，展現了出色的擴展性。隨著報告的發布，Sora 的核心研發人員 William Peebles 和紐約大學電腦科學助理教授謝賽寧合作撰寫的《Scalable Diffusion Models with Transformers》論文備受研究者關注。研究界希望透過論文中提出的 DiT 架構，探索再現 Sora 的可行性途徑。
最近，新加坡國立大學尤洋團隊開源的一個名為 OpenDiT 的專案為訓練和部署 DiT 模式開啟了新思路。
OpenDiT是一個專為提升DiT應用程式的訓練和推理效率而設計的系統，它不僅易於操作，而且速度快且記憶體利用高效。該系統涵蓋了文字到視訊生成和文字到圖像生成等功能，旨在為用戶提供高效、便利的體驗。
專案網址：https://github.com/NUS-HPC-AI-Lab/OpenDiT
OpenDiT 提供由Colossal-AI 支援的Diffusion Transformer (DiT) 的高效能實作。在訓練時，視訊和條件資訊分別被輸入到對應的編碼器中，作為DiT模型的輸入。隨後，透過擴散方法進行訓練和參數更新，最終將更新後的參數同步至EMA（Exponential Moving Average）模型。推理階段則直接使用EMA模型，將條件資訊作為輸入，從而產生對應的結果。
圖片來源：https://www.zhihu.com/people/berkeley-you-yang
OpenDiT 利用了ZeRO 平行策略，將DiT 模型參數分佈到多台機器上，初步降低了顯存壓力。為了取得更好的性能與精準度平衡，OpenDiT 也採用了混合精準度的訓練策略。具體而言，模型參數和優化器使用 float32 進行存儲，以確保更新的準確性。在模型計算的過程中，研究團隊為 DiT 模型設計了 float16 和 float32 的混合精度方法，以在維持模型精度的同時加速計算過程。
DiT 模型中使用的 EMA 方法是一種用於平滑模型參數更新的策略，可以有效提高模型的穩定性和泛化能力。但是會額外產生一份參數的拷貝，增加了顯存的負擔。為了進一步降低這部分顯存，研究團隊將 EMA 模型分片，並分別儲存在不同的 GPU 上。在訓練過程中，每個 GPU 只需計算和儲存自己負責的部分 EMA 模型參數，並在每次 step 後等待 ZeRO 完成更新後進行同步更新。
FastSeq
#在DiT 等視覺生成模型領域，序列並行性對於有效的長序列訓練和低延遲推理是必不可少的。
然而，DeepSpeed-Ulysses、Megatron-LM Sequence Parallelism 等現有方法在應用於此類任務時面臨局限性—— 要么是引入過多的序列通信，要么是在處理小規模序列並行時缺乏效率。
為此，研究團隊提出了 FastSeq，一種適用於大序列和小規模並行的新型序列並行。 FastSeq 透過為每個 transformer 層僅使用兩個通訊運算子來最小化序列通信，利用 AllGather 來提高通訊效率，並策略性地採用非同步 ring 將 AllGather 通訊與 qkv 計算重疊，進一步優化效能。
算符最佳化
在DiT 模型中引入adaLN 模組將條件資訊融入視覺內容，雖然這項操作對模型的性能提升至關重要，但也帶來了大量的逐元素操作，並且在模型中被頻繁調用，降低了整體的計算效率。為了解決這個問題，研究團隊提出了高效的 Fused adaLN Kernel，將多次操作合併成一次，從而增加了計算效率，並且減少了視覺資訊的 I/O 消耗。
圖片來源：https://www.zhihu.com/people/berkeley-you-yang
簡單來說，OpenDiT 有以下效能優勢：
#1、在GPU 上加速高達80%，50%的記憶體節省
2、FastSeq：一種新穎的序列平行方法
3、易於使用
4、文字到圖像和文字到影片產生完整pipeline
要使用OpenDiT，首先要安裝先決條件：
建議使用Anaconda 建立一個新環境（Python >= 3.10）來執行範例：
conda create -n opendit pythnotallow=3.10 -yconda activate opendit
#安裝ColossalAI：
git clone https://github.com/hpcaitech/ColossalAI.gitcd ColossalAIgit checkout adae123df3badfb15d044bd416f0cf29f250bc86pip install -e .
#安裝OpenDiT：
git clone https://github.com/oahzxl/OpenDiTcd OpenDiTpip install -e .
#（可選擇但推薦）安裝函式庫以加快訓練和推理速度：
# Install Triton for fused adaln kernelpip install triton# Install FlashAttentionpip install flash-attn# Install apex for fused layernorm kernelgit clone https://github.com/NVIDIA/apex.gitcd apexgit checkout 741bdf50825a97664db08574981962d66436d16apip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-optinotallow=--cpp_ext" --config-settings "--build-optinotallow=--cuda_ext" ./--global-optinotallow="--cuda_ext" --global-optinotallow="--cpp_ext"
#圖片生成
# Use scriptbash train_img.sh# Use command linetorchrun --standalone --nproc_per_node=2 train.py \--model DiT-XL/2 \--batch_size 2
enable_modulate_kernel: 是否啟用 modulate 核心最佳化，以加快訓練過程。預設值為 False，建議在 GPU
# Use scriptbash sample_img.sh# Use command linepython sample.py --model DiT-XL/2 --image_size 256 --ckpt ./model.pt
视频生成
你可以通过执行以下命令来训练视频 DiT 模型：
# train with sciptbash train_video.sh# train with command linetorchrun --standalone --nproc_per_node=2 train.py \--model vDiT-XL/222 \--use_video \--data_path ./videos/demo.csv \--batch_size 1 \--num_frames 16 \--image_size 256 \--frame_interval 3# preprocess# our code read video from csv as the demo shows# we provide a code to transfer ucf101 to csv formatpython preprocess.py
使用 DiT 模型执行视频推理的代码如下所示：
# Use scriptbash sample_video.sh# Use command linepython sample.py \--model vDiT-XL/222 \--use_video \--ckpt ckpt_path \--num_frames 16 \--image_size 256 \--frame_interval 3
为了验证 OpenDiT 的准确性，研究团队使用 OpenDiT 的 origin 方法对 DiT 进行了训练，在 ImageNet 上从头开始训练模型，在 8xA100 上执行 80k step。以下是经过训练的 DiT 生成的一些结果：
损失也与 DiT 论文中列出的结果一致：
要复现上述结果，需要更改 train_img.py 中的数据集并执行以下命令：
torchrun --standalone --nproc_per_node=8 train.py \--model DiT-XL/2 \--batch_size 180 \--enable_layernorm_kernel \--enable_flashattn \--mixed_precision fp16
感兴趣的读者可以查看项目主页，了解更多研究内容。
以上是想訓練類Sora模型嗎？尤洋團隊OpenDiT實現80%加速的詳細內容。更多資訊請關注PHP中文網其他相關文章！
