数日間GitHubのホットリストを独占した後、Colossal-AIが正式にリリースされました

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

超並列 AI トレーニングシステムである Colossal-AI は、ディープラーニングフレームワークの中核として機能するように設計されており、ユーザーが AI 展開の効率を簡単に最大化し、展開コストを最小限に抑えるのに役立ちます。

オープンソースアドレス: https://github.com/hpcaitech/ColossalAI

Colossal-AIはオープンソース化されて以来、広く注目を集め、GitHubのPythonホットリストで何日も連続して世界のトップを占めています。数万のスターを獲得した多くのスターオープンソースプロジェクトとともに、国内外で注目を集めています！

開発者の継続的な努力により、数か月に及ぶ集中的なテストを経て、ついに Colossal-AI が正式にリリースされました。このバージョンは 300 を超えるコミットで構成されています。

この公式バージョンアップデートは、分散トレーニングのパフォーマンスと開発者の使いやすさを最適化することに重点を置いています。主なハイライトは次のとおりです。

パフォーマンスと使いやすさを向上させるために ZeRO をリファクタリングしました。
トレーニング中にメモリ、ネットワーク、その他のステータスを監視するためのきめ細かい Profiler TensorBoard 監視プラグインを追加しました。
より柔軟なチェックポイント戦略とスケーラブルなパイプラインモジュール。
オープンソースのタンパク質予測 FastFold およびその他の業界ソリューション。
中国語のチュートリアル、MOE、BERT などの例を追加し、ユーザーコミュニティとフォーラムを開きます。

大規模モデルトレーニングのための専門的なサポート

近年、ディープラーニングの台頭と大規模モデルが主要なパフォーマンスチャートを席巻する中、最先端の AI モデルのサイズはわずか数年で 1 万倍に増加し、ハードウェアの緩やかな成長を数倍も上回っています。最先端の AI 大規模モデルは単一の GPU の容量をはるかに超えるだけでなく、必要な計算能力は単一の GPU で数百年、あるいは数千年にわたって実行する必要があることもよくあります。

そのため、単一の GPU の容量をどのように向上させるか、分散テクノロジを効率的に使用する方法、複数の GPU を組み合わせて低コストの並列トレーニング高速化を実現する方法が、大規模 AI モデルの主要な問題点となっています。

Colossal-AI は、並列次元の制限、効率の低さ、汎用性の低さ、導入の難しさ、メンテナンスの欠如など、既存のソリューションの問題点に対処するために、効率的な多次元並列処理、ビデオメモリの最適化、大規模な最適化ライブラリ、きめ細かな監視を使用して、ユーザーが最小限の変更で大規模な AI モデルトレーニングを効率的かつ迅速に導入できるようにします。

多次元並列処理

並列次元にデータ並列、1 次元テンソル並列、パイプライン並列のみが含まれる既存のソリューションと比較して、Colossal-AI は、2/2.5/3 次元テンソル並列とシーケンス並列、および便利な多次元ハイブリッド並列ソリューションも提供します。

△並列ViTテンソルの数が64の場合、バッチサイズは14倍に増加し、トレーニング速度は5倍に増加します。

その中で、高次元テンソル並列処理は、ビデオメモリの消費を大幅に削減し、通信効率を向上させ、コンピューティングリソースの利用をより効率的にすることができます。

△シーケンス並列処理により、BERT のトレーニング速度は 2 倍、シーケンス長は 1.5 倍に向上します。

シーケンス並列処理は、元のマシン機能の限界を打ち破り、大きな画像、ビデオ、長いテキスト、長期の医療モニタリングなどの長いシーケンスデータを直接処理するのに役立ちます。

ビデオメモリの最適化

Colossal-AI は、多次元並列処理、ZeRO 冗長メモリ除去、CPU オフロード、勾配チェックポイント、自動混合精度 (AMP) などの最先端テクノロジを含む複数のグラフィックスメモリ最適化テクノロジを統合し、ユーザーがグラフィックスメモリのボトルネックを回避し、トレーニングハードウェア要件を最大限に削減できるようにします。

△GPT-2はColossal-AIを使用しており、同じハードウェアでトレーニング可能なモデルサイズが24倍、トレーニング速度が3倍に増加します。

柔軟で使いやすい

Colossal-AI のインターフェース設計は PyTorch スタイルと一致しており、学習と使用のコストが削減されます。既存のプロジェクトは最小限の変更で Colossal-AI と組み合わせることができ、大規模な並列処理に簡単に拡張できます。さらに、システムは優れた拡張性を維持しており、必要に応じて新しい機能を簡単に追加でき、既存の機能モジュールとの互換性も確保されています。

きめ細かな監視: きめ細かな Profiler TensorBoard プラグインは、PyTorch と比較して、反復でのみトレーニングプロセスを記録できます。Colossal-AI は、反復内のネットワーク、通信、メモリなどの状態を監視できるため、開発者は正確な分析とデバッグを簡単に実行でき、開発効率が向上します。

大規模最適化ライブラリ: Colossal-AI は、LAMB や LARS などの大規模並列最適化ツールを提供し、初めてトレーニングバッチサイズを 65536 に拡張しました。 Colossal-AI は PyTorch に付属するさまざまなオプティマイザーとも互換性があり、さまざまなモデルのニーズを満たすために最新の最先端の最適化テクノロジーを常に探求し、追加しています。

豊富な業界ソリューション

Colossal-AIは現在、自動運転、クラウドコンピューティング、小売、医療、チップなどの業界の有名メーカーと協力関係を築いており、Hugging FaceなどのAI分野のトップオープンソース組織との協力関係も確立しています。

タンパク質構造予測加速ソリューション: FastFold

AlphaFoldは、タンパク質構造を予測する強力なAI能力により、ScienceとNatureによって2021年のトップ10の科学的ブレークスルーの1つに選ばれましたが、トレーニング時間が長い、コストが高いなどの問題があります。

△画像出典：https://arxiv.org/pdf/2203.00854.pdf

Colossal-AI に基づく高速化ソリューションである FastFold は、GPU 最適化と大規模モデルトレーニング技術を AlphaFold のトレーニングと推論に導入し、Google とコロンビア大学のソリューションを上回ることに成功しました。AlphaFold のトレーニング時間は 11 日から 67 時間に短縮され、総コストも削減され、長いシーケンスの推論では 9.3 倍から 11.6 倍の速度向上が実現しました。

△長系列推論性能の比較

GPUの半分はGPT-3をトレーニングする

GPT-3 などの非常に大規模な AI モデルの場合、Colossal-AI では、NVIDIA のソリューションと比較して、トレーニングを開始するために必要なコンピューティングリソースが半分しか必要ありません。同じコンピューティングリソースを使用した場合、速度は 11% 向上し、GPT-3 トレーニングのコストを 100 万ドル以上削減できます。

Colossal-AI はオープンソースコミュニティの構築に注力しており、中国語のチュートリアルを提供し、ユーザーコミュニティとフォーラムを開設し、ユーザーからのフィードバックに基づいて効率的なコミュニケーションと反復的な更新を実施し、MoE などの最先端のアプリケーションを継続的に追加しています。

プロジェクトチーム

Luchen の技術チームの中核メンバーは、カリフォルニア大学バークレー校、スタンフォード大学、清華大学、北京大学、シンガポール国立大学、シンガポール南洋理工大学など国内外の有名大学の出身者であり、Google Brain、IBM、Intel、Microsoft、NVIDIA などの有名メーカーでの勤務経験を持っています。同社は設立直後から、Innovation WorksやZhenFundなど複数のトップVC機関からシードラウンドの投資を受けた。

△ Luchen Technology 創設者 You Yang 教授：カリフォルニア大学バークレー校博士号、IPDPS/ICPP 最優秀論文賞、ACM/IEEE George Michael HPC フェローシップ、Forbes 30 Under 30 (Asia 2021)、IEEE-CS スーパーコンピューティング優秀新人賞、カリフォルニア大学バークレー校 EECS Lotfi A. Zadeh 優秀大学院賞

△ ルーチェンCSO ジェームズ・デメル教授：カリフォルニア大学バークレー校の特別教授、ACM/IEEEフェロー、米国科学アカデミー、米国工学アカデミー、米国芸術科学アカデミーの会員