ソラのようなモデルをトレーニングしたいですか? You Yang のチーム OpenDiT が 80% の高速化を達成

2024 年初頭のソラの驚異的なパフォーマンスは新たなベンチマークとなり、文生動画を研究しているすべての人々に追いつくために急ぐよう促しています。研究者は皆、ソラの結果を再現することに熱心で、時間との闘いに取り組んでいます。

OpenAI が公開した技術レポートによると、Sora の重要な革新ポイントは、ビジュアル データをパッチの統一表現に変換し、それを Transformer と拡散モデルと組み合わせて実証することです。優れたパフォーマンス、拡張性。このレポートの発表に伴い、Sora の中心開発者である William Peebles とニューヨーク大学コンピューター サイエンス助教授 Xie Saining が共著した論文「Scalable Diffusion Models with Transformers」は研究者から大きな注目を集めています。研究コミュニティは、論文で提案されている DiT アーキテクチャを通じてソラを再現する実現可能な方法を探ることを望んでいます。

最近、シンガポール国立大学の You Yang チームがオープンソース化した OpenDiT と呼ばれるプロジェクトにより、DiT モデルのトレーニングとデプロイのための新しいアイデアが開かれました。

OpenDiT は、DiT アプリケーションのトレーニングと推論の効率を向上させるために設計されたシステムで、操作が簡単なだけでなく、高速でメモリ効率も優れています。このシステムは、テキストからビデオへの生成やテキストから画像への生成などの機能をカバーしており、ユーザーに効率的で便利なエクスペリエンスを提供することを目指しています。

プロジェクトアドレス: https://github.com/NUS-HPC-AI-Lab/OpenDiT

OpenDiT メソッドの紹介

OpenDiT は、Colossal-AI を活用した拡散トランス (DiT) の高性能実装を提供します。トレーニング中に、ビデオと状態の情報が、DiT モデルへの入力としてそれぞれ対応するエンコーダーに入力されます。その後、拡散法によってトレーニングとパラメータの更新が実行され、最後に更新されたパラメータが EMA (指数移動平均) モデルに同期されます。推論段階では、EMA モデルが直接使用され、条件情報を入力として受け取り、対応する結果を生成します。

画像出典: https://www.zhihu.com/people/berkeley-you-yang

OpenDiT は、ZeRO 並列戦略を使用して DiT モデル パラメーターを複数のマシンに分散し、最初にメモリ負荷を軽減します。パフォーマンスと精度のより良いバランスを達成するために、OpenDiT は混合精度トレーニング戦略も採用しています。具体的には、正確な更新を保証するために、モデル パラメーターとオプティマイザーは float32 を使用して保存されます。モデルの計算プロセス中に、研究チームは、モデルの精度を維持しながら計算プロセスを高速化するために、DiT モデル用に float16 と float32 の混合精度手法を設計しました。

DiT モデルで使用される EMA メソッドは、モデル パラメーターの更新をスムーズにするための戦略であり、モデルの安定性と汎化能力を効果的に向上させることができます。ただし、パラメータの追加コピーが生成されるため、ビデオ メモリへの負担が増加します。ビデオ メモリのこの部分をさらに削減するために、研究チームは EMA モデルを断片化し、異なる GPU に保存しました。トレーニング プロセス中、各 GPU は EMA モデル パラメーターの独自の部分を計算して保存し、同期更新の各ステップの後に ZeRO が更新を完了するのを待つだけで済みます。

FastSeq

DiT などの視覚生成モデルの分野では、効率的な長いシーケンスを実現するためにシーケンスの並列性が重要です。訓練と遅延の少ない推論が不可欠です。

しかし、DeepSpeed-Ulysses、Megatron-LM シーケンス並列処理などの既存の手法は、そのようなタスクに適用すると制限に直面します。つまり、大量のシーケンス通信が導入されたり、小規模な処理を扱う際に効率が欠如したりすることになります。 -スケール順次並列処理。

この目的を達成するために、研究チームは、大規模なシーケンスと小規模な並列処理に適した新しいシーケンス並列処理である FastSeq を提案しました。 FastSeq は、トランス層あたり 2 つの通信オペレーターのみを使用することでシーケンス通信を最小限に抑え、AllGather を活用して通信効率を向上させ、非同期リングを戦略的に採用して AllGather 通信を qkv 計算とオーバーラップさせてパフォーマンスをさらに最適化します。

オペレータの最適化# ＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃＃

adaLN モジュールは、条件付き情報をビジュアル コンテンツに統合するために DiT モデルに導入されています。この操作はモデルのパフォーマンスを向上させるために重要ですが、多数の要素ごとの操作も必要になります。モデル内で頻繁に使用される呼び出しにより、全体的なコンピューティング効率が低下します。この問題を解決するために、研究チームは、複数の演算を 1 つに結合する効率的な Fused adaLN カーネルを提案しました。これにより、計算効率が向上し、視覚情報の I/O 消費が削減されます。

conda create -n opendit pythnotallow=3.10 -yconda activate opendit

git clone https://github.com/hpcaitech/ColossalAI.gitcd ColossalAIgit checkout adae123df3badfb15d044bd416f0cf29f250bc86pip install -e .

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

# 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

# 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

torchrun --standalone --nproc_per_node=8 train.py \--model DiT-XL/2 \--batch_size 180 \--enable_layernorm_kernel \--enable_flashattn \--mixed_precision fp16