ColossalAI: 大規模なモデルの効率的なトレーニングと推論のための分散ディープラーニングフレームワーク

解決する課題

Colossal-AIは、大規模なAIモデルのトレーニングと推論を、より安価に、より速く、よりアクセスしやすくすることを目的に設計されています。複数のGPUやハードウェア構成にわたってディープラーニングモデルをスケールアップする際に伴う、高い計算コストとメモリの制限という課題に対処します。

仕組み

このプロジェクトは、開発者が単一のノートパソコンで行うのと同じくらい簡単に分散ディープラーニングモデルを記述できるようにする、一連の並列コンポーネントとメモリ管理ツールを提供します。以下のような、いくつかの高度な並列化戦略を採用しています：

• 並列化戦略: Data Parallelism、Pipeline Parallelism、および様々な形式のTensor Parallelism (1D, 2D, 2.5D, 3D)、ならびにSequence ParallelismとZero Redundancy Optimizer (ZeRO)。
• Auto-Parallelism: モデルのハードウェアへの分散を自動的に処理します。
• Heterogeneous Memory Management: PatrickStarのようなツールを使用して、異なるハードウェア層にわたってメモリを管理します。
• Configuration-based Usage: 設定ファイルを通じて並列化設定を定義できるため、ユーザーフレンドリーな体験を提供します。

対象ユーザー

LLM、Soraのような動画生成モデル、またはStable Diffusionのような画像生成モデルなどのモデルを、大規模なクラスターにスケールアップしたり、コンシューマーグレードのGPU向けに最適化したりする必要があるAI研究者や開発者を対象としています。

ハイライト

• 幅広いモデルサポート: LLaMA 1/2/3、GPT-3、BERT、PaLM、およびMoEモデル向けの最適化された実装を含みます。
• 大幅なパフォーマンス向上: ベンチマークでは、B200やH200のようなハイエンドGPUにおいて大幅なスループットの向上が示されています。
• 実世界でのアプリケーション: 動画生成のためのOpen-Soraや、完全なRLHFパイプラインを備えたChatGPTのクローンを作成するためのColossalChatなどのプロジェクトを支えています。
• メモリ効率: Stable Diffusionのトレーニングにおけるメモリ消費量を最大5.6倍削減することができ、RTX 3060のような低スペックのハードウェアでのトレーニングを可能にします。