ColossalAI: 効率的な大規模AIモデル学習ソリューションの提供

はじめに

ColossalAIは、HPC-AI Technologies社によって開発されたオープンソースプラットフォームで、大規模なAIモデルの学習と推論に効率的で費用対効果の高いソリューションを提供します。複数の並列戦略、ヘテロジニアスメモリ管理、混合精度学習をサポートすることにより、ColossalAIはモデルの学習と推論にかかる時間とリソースの消費を大幅に削減することができます。データ並列、テンソル並列、パイプライン並列のいずれにおいても、ColossalAIは、研究者や開発者がマルチGPUクラスタ上で大規模モデルの効率的な学習と推論を実現するための強力なツールとライブラリを提供します。

機能一覧

• データ並列、テンソル並列、パイプライン並列、その他の並列ストラテジーのサポート
• ミックスド・プレシジョン・トレーニングとゼロ・リダンダンシー・オプティマイザー（ZeRO）
• 大規模モデルの効率的な学習をサポートするヘテロジニアスメモリ管理
• Open-Sora、Colossal-LaMAなど、複数のドメイン固有モデルをサポート。
• 分散学習と推論のための使いやすいツールの提供
• 高性能カーネル、KVキャッシュ、ページング・アテンション、シーケンシャル・バッチ処理の統合
• 設定ファイルによる並列トレーニングの容易な設定
• 豊富な例とドキュメントで、すぐに使い始められる
• Dockerイメージとソースコードからのビルドのための複数のインストールオプションを提供します。

ヘルプの使用

インストレーション・ガイド

PyPIからのインストール

Colossal-AIは以下のコマンドで簡単にインストールできます：

pip install colossalai

デフォルトでは、PyTorch の拡張機能はインストール中にビルドされません。PyTorch の拡張機能をビルドする必要がある場合はBUILD_EXT=1::

BUILD_EXT=1 pip install colossalai

さらに、毎週ナイトバージョンをリリースしており、未発表の最新機能やバグ修正にアクセスすることができます。インストール方法は以下の通りです：

pip install colossalai-nightly

ソースからのインストール

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
pip install .

CUDA/C++カーネルはデフォルトではコンパイルされません。colossalAIは実行時にビルドします。必要であればCUDAカーネルフュージョンを有効にしてください：

BUILD_EXT=1 pip install .

CUDA 10.2ユーザーの場合、インストールする前に手動でcubライブラリをダウンロードし、適切なディレクトリにコピーすることができます。

Dockerの使用

DockerHubからイメージを取り込む

から直接情報を得ることができる。DockerHubのページDockerイメージを取り出します。

自分のイメージを作り上げる

cd ColossalAI
docker build -t colossalai ./docker

コンテナを対話モードで起動する：

docker run -ti --gpus all --rm --ipc=host colossalai bash

機能 操作の流れ

データ並列性

データ並列化とは、データセットを複数のサブセットに分割し、複数のGPU上で並列にモデルをトレーニングするプロセスです。ColossalAIは、簡素化されたデータ並列化プロファイルにより、ユーザーが簡単にデータ並列トレーニングを行えるようにします：

from colossalai.nn.parallel import DataParallel
model = DataParallel(model)

テンソル並列

テンソル並列は、モデルのパラメータテンソルを複数のサブテンソルに分割し、複数のGPU上で並列計算するプロセスです。ColossalAIは、1D、2D、2.5D、3Dのテンソル並列の実装を提供します：

from colossalai.nn.parallel import TensorParallel
model = TensorParallel(model, parallel_mode='1D')

組立ラインと平行に走る

パイプライン並列とは、モデルを複数のステージに分割し、それぞれを1つまたは複数のGPUで実行することです：

from colossalai.pipeline.parallel import PipelineParallel
model = PipelineParallel(model, num_stages=4)

ミックス精密トレーニング

混合精度トレーニングは、トレーニング中に16ビット浮動小数点数（FP16）と32ビット浮動小数点数（FP32）を組み合わせて使用することにより、メモリ使用量を大幅に削減し、トレーニングを高速化します：

from colossalai.amp import convert_to_amp
model, optimizer, criterion = convert_to_amp(model, optimizer, criterion)

ゼロ・リダンダンシー・オプティマイザー（ZeRO）

ZeROオプティマイザーは、オプティマイザーの状態、勾配、パラメータを複数のGPUに分散させることで、グラフィックス・メモリのフットプリントを大幅に削減します：

from colossalai.zero import ZeroOptimizer
optimizer = ZeroOptimizer(optimizer, model)

実世界での応用

オープン・ソラ

Open-Soraは、モデルパラメータ、トレーニングの詳細、ワンクリックで16秒の720p HDビデオを生成する機能を含む、ビデオ生成モデルのためのColossalAIの完全なソリューションです：

# 训练
python train.py
# 推理
python infer.py

詳細はこちらをご覧ください。オープン・ソラ.

コロッサル-LLaMA

Colossal-LaMAは、ドメインに特化した大規模言語モデル（LLM）のためのオープンソース・ソリューションを提供します：

# 训练
python train_llama.py
# 推理
python infer_llama.py

詳細はこちらをご覧ください。コロッサル-LLaMA.