ソラのようなモデルをトレーニングしたいですか? 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 消費が削減されます。

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

簡単に言うと、OpenDiT には次のようなパフォーマンス上の利点があります:

#1. GPU で最大 80% の高速化、50% のメモリ節約

• DiT 用に設計された Fused AdaLN をはじめ、FlashAttendant、Fused Layernorm、HybridAdam など、効率的なオペレーターを設計しました。
• ZeRO、Gemini、DDP などのハイブリッド並列アプローチを使用します。 ema モデルをシャーディングすると、メモリ コストもさらに削減されます。

#2. FastSeq: 新しいシーケンス並列メソッド

• は、同様の目的で設計されています。 DiT は、LLM に比べてシーケンスが通常長いもののパラメータが小さいワークロード向けに設計されています。
• ノード内シーケンス並列処理により、トラフィックを最大 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 -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 で使用されるブースター プラグインをサポートします。デフォルトは zero2 です。zero2 を有効にすることをお勧めします。
• mixed_precision: 混合精度トレーニングのデータ型。デフォルトは fp16 です。
• grad_checkpoint: グラデーション チェックポイントを有効にするかどうか。これにより、トレーニング プロセスのメモリ コストが節約されます。デフォルト値は False です。十分なメモリがある場合は、無効にすることをお勧めします。
• enable_modulate_kernel: トレーニング プロセスを高速化するためにカーネル最適化の調整を有効にするかどうか。デフォルト値は False で、H100 未満の GPU ではこれを有効にすることをお勧めします。
• enable_layernorm_kernel: トレーニング プロセスを高速化するために、layernorm カーネルの最適化を有効にするかどうか。デフォルト値は False で、これを有効にすることをお勧めします。
• enable_flashattn: トレーニング プロセスを高速化するために FlashAttendant を有効にするかどうか。デフォルト値は False で、これを有効にすることをお勧めします。
• sequence_Parallel_size: シ​​ーケンスの並列処理のサイズ。値 > 1 を設定すると、シーケンスの並列処理が有効になります。デフォルト値は 1 ですが、十分なメモリがある場合は無効にすることをお勧めします。

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

うわー

视频生成

你可以通过执行以下命令来训练视频 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

感兴趣的读者可以查看项目主页，了解更多研究内容。

以上がソラのようなモデルをトレーニングしたいですか? You Yang のチーム OpenDiT が 80% の高速化を達成の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。