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

2024年初頭のキング爆弾として、ソラの出現は追いつくための新たな目標を設定しました。ヴィンセントビデオのすべての研究者は、ソラの効果を最短時間で再現したいと考えています。

OpenAIが公開した技術レポートによると、Soraの核心技術ポイントの一つは、視覚データをパッチの統一された表現形式に変換することであり、Transformerと拡散モデルの組み合わせにより、優れたスケール特性を発揮するという。報告書が発表された後、SoraのコアR＆Dメンバーであるウィリアム・ピーブルズ氏とニューヨーク大学のコンピュータサイエンスの助教授である謝彩寧氏が共著した論文「トランスフォーマーによるスケーラブルな拡散モデル」が多くの研究者の注目を集めました。私たちは、この論文で提案された DiT アーキテクチャを突破口として、Sora を再現するための実現可能な方法を模索したいと考えています。

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

OpenDiT は、テキストからビデオへの生成やテキストから画像への生成を含む DiT アプリケーションのトレーニングと推論の効率を向上させるために特別に設計された、使いやすく、高速で、メモリ効率に優れたシステムです。

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

OpenDiTメソッドの紹介

OpenDiT は、Colossal-AI を搭載した Diffusion Transformer (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 が更新を完了するのを待って同期更新を実行します。

ファストシーケンス

DiT などの視覚生成モデルの分野では、長いシーケンスの効率的なトレーニングと低レイテンシの推論にはシーケンスの並列処理が不可欠です。

しかし、DeepSpeed-Ulysses や Megatron-LM Sequence Parallelism などの既存の方法をこのようなタスクに適用すると、シーケンス通信が多すぎるか、小規模なシーケンス並列処理の処理効率が不足するかのいずれかの制限に直面します。

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

オペレータの最適化

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

画像出典: 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: 新しいシーケンス並列アプローチ

• シーケンスは通常長いが、パラメーターは LLM に比べて小さい DiT のようなワークロード向けに設計されています。
• ノード内順次並列処理により、通信を最大 48% 節約できます。
• 単一の GPU のメモリ制限を打ち破り、全体的なトレーニングと推論の時間を短縮します。

3. 使いやすい

• ほんの数行のコードを変更するだけで、パフォーマンスを大幅に向上させることができます。
• ユーザーは分散トレーニングがどのように実装されるかを理解する必要はありません。

4. テキストから画像、テキストからビデオを生成する完全なパイプライン

• 研究者やエンジニアは、OpenDiT パイプラインを簡単に使用し、並列部分を変更することなく実際のアプリケーションに適用できます。
• 研究チームは、ImageNet でテキストを画像に変換するトレーニングを行うことで OpenDiT の精度を検証し、チェックポイントを公開しました。

インストールと使用方法

OpenDiT を使用するには、まず前提条件をインストールします。

• Python >= 3.10
• PyTorch >= 1.13 (バージョン > 2.0 を推奨)
• CUDA >= 11.6

例を実行するには、Anaconda を使用して新しい環境 (Python >= 3.10) を作成することをお勧めします。

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

ColossalAIをインストールします:

 git clone https://github.com/hpcaitech/ColossalAI.gitcd ColossalAI git checkout adae123df3badfb15d044bd416f0cf29f250bc86 pip install -e .

OpenDiT をインストールします。

 git clone https://github.com/oahzxl/OpenDiTcd OpenDiT pip install -e .

(オプションですが推奨) トレーニングと推論を高速化するためのライブラリをインストールします。

 # Install Triton for fused adaln kernel pip install triton # Install FlashAttention pip install flash-attn # Install apex for fused layernorm kernel git clone https://github.com/NVIDIA/apex.gitcd apex git checkout 741bdf50825a97664db08574981962d66436d16a pip 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 script bash train_img.sh# Use command line torchrun --standalone --nproc_per_node=2 train.py \ --model DiT-XL/2 \ --batch_size 2

デフォルトでは、すべての加速方法は無効になっています。トレーニング プロセスの重要な要素の詳細は次のとおりです。

• プラグイン: ColossalAI、zero2、ddp で使用されるブースター プラグインをサポートします。デフォルトは zero2 ですが、zero2 を有効にすることをお勧めします。
• mixed_precision: 混合精度トレーニングのデータ型。デフォルトは fp16 です。
• grad_checkpoint: 勾配チェックポイントを有効にするかどうか。これにより、トレーニング中のメモリコストが節約されます。デフォルト値は False です。メモリが十分にある場合は無効にすることをお勧めします。
• enable_modulate_kernel: トレーニング プロセスを高速化するために、モジュレート カーネルの最適化を有効にするかどうか。デフォルト値は False です。GPU < H100 の場合は有効にすることをお勧めします。
• enable_layernorm_kernel: トレーニング プロセスを高速化するために、layernorm カーネル最適化を有効にするかどうか。デフォルト値は False ですが、有効にすることをお勧めします。
• enable_flashatn: トレーニング プロセスを高速化するために FlashAttention を有効にするかどうか。デフォルト値は False ですが、有効にすることをお勧めします。
• sequence_parallel_size: シ​​ーケンスの並列サイズ。これを 1 より大きい値に設定すると、シーケンスの並列処理が有効になります。デフォルト値は 1 ですが、メモリが十分にある場合は無効にすることをお勧めします。

推論に DiT モデルを使用する場合は、次のコードを実行し、チェックポイント パスを独自のトレーニング済みモデルに置き換えることができます。

 # Use script bash sample_img.sh# Use command line python sample.py --model DiT-XL/2 --image_size 256 --ckpt ./model.pt

ビデオ生成

次のコマンドを実行して、ビデオ DiT モデルをトレーニングできます。

 # train with scipt bash train_video.sh# train with command line torchrun --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 format python preprocess.py

DiT モデルを使用してビデオ推論を実行するコードは次のとおりです。

 # Use script bash sample_video.sh# Use command line python 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 で 80,000 ステップを実行しました。トレーニングされた 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

興味のある読者はプロジェクトのホームページにアクセスして、研究の詳細を知ることができます。