想訓練類Sora模型嗎？尤洋團隊OpenDiT實現80%加速

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 方法介紹

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%的記憶體節省

• #設計了高效率的算子，包括針對DiT設計的Fused AdaLN，以及FlashAttention、Fused Layernorm 和HybridAdam。
• 採用混合平行方法，包括 ZeRO、Gemini 和 DDP。對 ema 模型進行分片也進一步降低了記憶體成本。

2、FastSeq：一種新穎的序列平行方法

• 專為類似DiT 的工作負載而設計，在這些應用中，序列通常較長，但參數與LLM 相比較小。
• 節點內序列並行可節省高達 48% 的通訊量。
• 打破單一 GPU 的記憶體限制，減少整體訓練和推理時間。

3、易於使用

• #只需幾行程式碼的修改，即可獲得巨大的性能提升。
• 使用者無需了解分散式訓練的實作方式。

4、文字到圖像和文字到影片產生完整pipeline

• 研究人員和工程師可以輕鬆使用OpenDiT pipeline 並將其應用於實際應用，而無需修改平行部分。
• 研究團隊透過在 ImageNet 上進行文字到圖像訓練來驗證 OpenDiT 的準確性，並發布了檢查點（checkpoint）。

安裝與使用

要使用OpenDiT，首先要安裝先決條件：

• Python >= 3.10
• PyTorch >= 1.13（建議使用>2.0 版本）
• CUDA > = 11.6

建議使用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"

#圖片生成

• 你可以執行以下指令來訓練DiT 模型：

• # Use scriptbash train_img.sh# Use command linetorchrun --standalone --nproc_per_node=2 train.py \--model DiT-XL/2 \--batch_size 2

• 預設禁用所有加速方法。以下是訓練過程中一些關鍵要素的詳細資訊：
• plugin: 支援 ColossalAI、zero2 和 ddp 使用的 booster 外掛程式。預設是 zero2，建議啟用 zero2。
• mixed_ precision：混合精準度訓練的資料類型，預設是 fp16。
• grad_checkpoint: 是否啟用梯度檢查點。這節省了訓練過程的記憶體成本。預設值為 False。建議在內存足夠的情況下禁用它。

enable_modulate_kernel: 是否啟用 modulate 核心最佳化，以加快訓練過程。預設值為 False，建議在 GPU

enable_layernorm_kernel: 是否啟用 layernorm 核心最佳化，以加快訓練過程。預設值為 False，建議啟用它。

#########enable_flashattn: 是否啟用 FlashAttention，以加快訓練過程。預設值為 False，建議啟用。 ############sequence_parallel_size：序列並行度大小。當設定值 > 1 時將啟用序列並行。預設值為 1，如果記憶體足夠，建議停用它。 ###############如果你想使用 DiT 模型進行推理，可以運行以下程式碼，需要將檢查點路徑替換為你自己訓練的模型。 #########

# 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

DiT 复现结果

为了验证 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中文網其他相關文章！