1 つの GPU + 数行のコードで、大規模モデルのトレーニング速度が 40% 向上します。 HuggingFaceのシームレスなサポート

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

大型モデルをより多くの人に使ってもらうために、技術サークルではさまざまなユニークなアイデアを考案してきたと言わざるを得ません。

モデルが十分に開いていませんか?無料のオープンソースバージョンを独自に開発した人もいます。

例えば、最近ネット上で話題になっているDALL・E Miniや、Metaが公開したOPT-175B （Open Pretrained Transformer）などです。

レプリケーションという方法により、これまで十分に公開されていなかった大規模なモデルを誰でも利用できるようになります。

一方、モデルが大きすぎるため、個々のプレーヤーが莫大なコストを負担するのは困難だと考える人もいます。

そのため、大規模モデルのトレーニングを高速化し、コストを削減するために、異種メモリや並列コンピューティングなどの方法が提案されています。

たとえば、オープンソースプロジェクトColossal-AI は最近、Nvidia 3090 が 180 億のパラメータを持つ大規模モデルに単独で挑戦できるようにするという目標を達成しました。

そして、過去 2 日間で、同社は新たな一連の新製品を発売しました。

Hugging Face コミュニティモデルをシームレスにサポートし、わずか数行のコードを追加するだけで、低コストで大規模なモデルのトレーニングと微調整が可能になります。

Hugging Face は現在最も人気のある AI ライブラリの 1 つであり、 50,000 を超えるAI モデルの実装を提供し、多くの AI プレーヤーが大規模なモデルをトレーニングするための第一の選択肢となっていることを知っておく必要があります。

Colossal-AI のこの動きにより、公開モデルのトレーニングと微調整がより実現可能になります。

トレーニング効果も向上します。

単一の GPUでは、Microsoft の DeepSpeed と比較して、Colossal-AI の自動最適化戦略を使用することで、最大40%の高速化を実現できます。

しかし、PyTorch などの従来のディープラーニングフレームワークでは、このような大規模なモデルを単一の GPU で実行できなくなりました。

8 つの GPU を使用して並列トレーニングを行う場合は、起動コマンドに-nprocs 8 を追加するだけです。

この進歩の波は、個々の AI プレーヤーが考慮する必要があるコスト、効率、および実用的な問題を把握したと言えます。

コードロジックを変更する必要はありません

何もせずに話すのは、ただの空論に過ぎません。

OPT を例に、Colossal-AI の新機能の使い方を詳しく見てみましょう。

OPT の正式名称は Open Pretrained Transformer です。

これは Meta AI によってリリースされたもので、最大パラメータサイズが 1750 億でGPT-3 に匹敵します。

最大の特徴は、GPT-3 はモデルの重みを公開しないのに対し、OPT はすべてのコードと重みをオープンソース化していることです。

したがって、すべての開発者はこれに基づいてパーソナライズされたダウンストリームタスクを開発できます。

次の例では、OPT によって提供される事前トレーニング済みの重みに基づいて因果言語モデルを微調整します。

主な手順は 2 つあります。

プロフィールを追加する
実行開始

最初のステップは、実行するタスクに応じて構成ファイルを追加することです。

たとえば、GPU で異種トレーニングを例にとると、構成ファイルに関連する構成項目を追加するだけでよく、コードのトレーニングロジックを変更する必要はありません。

たとえば、 tensor_placement_policy は異種トレーニングの戦略を決定し、そのパラメータは CUDA、CPU、および auto になります。

それぞれの戦略には異なる利点があり、さまざまな状況に適しています。

CUDA : すべてのモデルパラメータを GPU に配置します。これは、オフロードせずにトレーニングを実行できる従来のシナリオに適しています。

CPU : すべてのモデルパラメータはCPUメモリに配置され、現在計算に関係する重みのみがGPUメモリに保持されます。これは、非常に大規模なモデルのトレーニングに適しています。

auto : リアルタイムのメモリ情報に基づいて GPU メモリに保持するパラメータの量を自動的に決定し、CPU と GPU 間のデータ転送を削減しながら GPU メモリの使用を最大化します。

一般ユーザーにとっては、自動戦略を使用するのが最も便利です。

このようにして、Colossal-AI は、コンピューティング効率を最大化するために、リアルタイムで最適な異種戦略を自動的かつ動的に選択できます。

 colossalai.zero.shard_utils からTensorShardStrategy をインポートします


ゼロ= dict ( model_config = dict ( shard_strategy = TensorShardStrategy ()、
 tensor_placement_policy = "auto" )、
 optimizer_config = dict ( gpu_margin_mem_ratio = 0.8 )

2 番目の手順は、構成ファイルの準備ができたら、新しい機能を有効にするために数行のコードを挿入することです。

まず、1 行のコードを使用して構成ファイルで Colossal-AI を起動します。

Colossal-AI は、分散環境を自動的に初期化し、関連する構成を読み取り、構成内の関数をモデルやオプティマイザーなどのコンポーネントに自動的に挿入します。

 colossalai . launch_from_torch ( config = './configs/colossalai_zero.py' )

次に、データセット、モデル、オプティマイザー、損失関数などを通常どおり定義します。

たとえば、ネイティブ PyTorch コードを直接使用する場合、モデルを定義するときにZeroInitContextでモデルを初期化するだけで済みます。

ここでは、Hugging Face が提供する OPTForCausalLM モデルと事前トレーニング済みの重みが使用され、Wikitext データセットで微調整されています。

 ZeroInitContext ( target_device = torch.cuda.current_device () の場合、
 シャード戦略= シャード戦略、
 shard_param = True ):
 モデル= OPTForCausalLM . from_pretrained (
 'facebook/opt-1.3b'
 設定= 設定
 ）

次に、 colossalai.initializeを呼び出して、構成ファイルで定義された異種メモリ関数をトレーニングエンジンに挿入し、対応する関数を開始するだけです。

 エンジン、 train_dataloader 、 eval_dataloader 、 lr_scheduler = colossalai.initialize ( モデル= model 、
 オプティマイザー= オプティマイザー、
 基準= 基準、
 トレインデータローダー= トレインデータローダー、
 test_dataloader = eval_dataloader 、
 lr_scheduler = lr_scheduler )

依然としてGPU + CPUの異種性に頼る必要がある

ユーザーが上記の「絶対確実な」操作を実行できるようにするための鍵は、AI システム自体が十分にスマートである必要があるということです。

中核的な役割を担うのは、Colossal-AI システムの効率的な異種メモリ管理サブシステムであるGeminiです。

システム内の総支配人のような存在で、計算に必要な情報を収集した上で、CPUやGPUのメモリ使用量を動的に割り当てます。

具体的な動作原理は、最初の数ステップで予熱し、PyTorch の動的計算グラフでメモリ消費情報を収集することです。

ウォームアップが終了し、演算子を計算する前に、Gemini は収集されたメモリ使用量記録を使用して、コンピューティングデバイス上で演算子に必要なピークメモリを予約し、同時に一部のモデルテンソルを GPU ビデオメモリから CPU メモリに移動します。

Gemini の組み込みメモリマネージャーは、各テンソルに HOLD、COMPUTE、FREE などのステータスメッセージをマークします。

そして、動的に照会されたメモリ使用量に応じて、テンソルの状態が継続的かつ動的に変換され、テンソルの位置が調整されます。

直接的な利点は、ハードウェアが非常に限られている場合に、モデル容量を最大化し、トレーニング速度のバランスをとることができることです。

業界で主流の方式である ZeRO (Zero Reduency Optimizer) も CPU + GPU 異種メモリ方式を使用していますが、静的パーティショニングであるため、システムクラッシュ、不要な通信などの問題が発生する可能性があることに注意してください。

さらに、ダイナミックヘテロジニアスCPU+GPUメモリ方式を採用し、メモリスティックを追加することでメモリを拡張することもできます。

ハイエンドのグラフィックカードを購入するよりもはるかに安価です。

現在、Colossal-AI メソッドを使用すると、RTX 2060 6GB の一般的なゲームマシンは15 億のパラメータモデルをトレーニングできます。RTX 3090 24GB ホストは180 億のパラメータの大規模モデルに直接挑戦できます。Tesla V100 32GB は240 億のパラメータを処理することもできます。

Colossal-AI は、メモリの使用を最大化するだけでなく、分散並列方式を使用してトレーニング速度を継続的に向上させます。

データ並列処理、パイプライン並列処理、2.5 次元テンソル並列処理などの複雑な並列戦略を同時に使用することを提案します。

この方法は複雑ですが、使い方は非常に簡単です。簡単なステートメントだけで自動的に実装できます。

他のシステムやフレームワークのようにコードに介入して、複雑な基礎ロジックを手動で処理する必要はありません。

 並列= 辞書(
 パイプライン= 2 、
 テンソル= dict ( モード= '2.5d' 、 深さ= 1 、 サイズ= 4 )
 ）

Colossal-AI は他に何ができるのでしょうか?

実際、Colossal-AI はオープンソースとしてリリースされて以来、GitHub や Papers With Code のホットリストで何度も世界第 1 位にランクされており、テクノロジー界ではよく知られる存在となっています。

前述のように単一の GPU で大規模なモデルをトレーニングすることに加えて、Colossal-AI を数十または数百の GPU を使用した大規模な並列シナリオに拡張すると、パフォーマンスが 2 倍になり、リソース使用量が NVIDIA Megatron-LM などの既存のシステムの 10 分の 1 未満に削減されます。

つまり、GPT-3 などの大規模な AI モデルの事前トレーニングにかかるコスト削減は数百万元に達する可能性があります。

Colossal-AI関連のソリューションは、自動運転、クラウドコンピューティング、小売、医療、チップなどの業界の有名メーカーによって使用されていることが明らかになっています。

同時に、オープンソースコミュニティの構築にも大きな注意を払い、中国語のチュートリアルを提供し、ユーザーコミュニティフォーラムをオープンし、全員のニーズとフィードバックに基づいて継続的に更新と反復を行っています。

たとえば、Colossal-AI が Hugging Face にいくつかのモデルを直接ロードできるかどうかを尋ねるメッセージを残しているファンがいることがわかりました。

さて、アップデートが来ました。

では、大規模モデルのトレーニングでは、現在どのような困難を解決する必要があると思いますか?

コメント欄にメッセージを残して議論しましょう〜