大規模なモデルをトレーニングするのは本当に難しいのでしょうか?事前トレーニング済みで使いやすく、非常に効率的な「Li Bai」モデルライブラリが登場しました。

LiBai モデルライブラリは、Hugging Face、Megatron-LM、DeepSpeed、FairSeq などの主流の Transformer ライブラリの利点をすべて網羅しており、一般の人々が大規模なモデルトレーニングを利用できるようになります。

大きなモデルがたくさんあるのですが、速度を上げる方法を教えてください。 2018年にBERTが誕生し、GPT-3やViTなど数億のパラメータを持つモデルが登場して以来、AIモデルのパラメータが爆発的に増加することはもはや驚くべきことではなく、錬金術師たちは注意を払う暇もなく、麻痺状態に陥ることさえあります。

同時に、大規模なモデルはコンピューティングとメモリのリソースに大きな課題をもたらします。トレーニングのコストは急激に上昇しています。たとえば、非常に高度な NVIDIA A100 GPU を使用して、数千億のパラメータを持つモデル GPT-3 をトレーニングするには 100 年以上かかります。

大規模モデル用のビデオメモリの需要は、GPU ビデオメモリの成長率よりもはるかに速いペースで増加しています。OpenAI のレポートによると、モデルサイズは 3.5 か月ごとに 2 倍になりますが、GPU ビデオメモリが 2 倍になるには 18 か月かかります。 GPU メモリの制限により、単一の GPU では大規模なモデルパラメータに対応できなくなります。

そのため、業界ではコンピューティングを複数の GPU デバイスに拡張する必要があり、大多数の開発者にとって分散トレーニングは避けられない選択肢となっています。

しかし、分散トレーニングの敷居は高すぎます。十分なコンピューティングリソースがあっても、分散トレーニングを処理できないためにイライラしてしまう可能性があります。分散並列プログラミングでは通常、エンジニアがコンピューターシステムとアーキテクチャに関する専門知識と関連する実務経験を持つ必要があるため、最先端のアルゴリズムと新しいモデルを探索することがさらに困難になります。これらすべての要因により、大規模モデルは一部のテクノロジー大手の特権となっています。モデルのトレーニング効率を向上し、より多くのエンジニアが大規模なモデルを使用および研究できるようにすることが最優先事項となっています。

問題は、分散トレーニングをサポートするモデルライブラリが市場に多数存在する中で、どれが最も適しているかということです。

最近、一流の技術チームが開発した高効率の国産オープンソースディープラーニングフレームワークOneFlowが、LiBaiモデルライブラリを発表しました。この新世代モデルライブラリは、Hugging Face、Megatron-LM、DeepSpeed、FairSeqなど、主流のTransformerライブラリの利点をすべて網羅しており、分散トレーニングのパフォーマンスは相変わらず優れています。さらに重要なのは、グローバルパースペクティブプログラミングにより、分散使用の敷居が最大限に下がり、一般の人々が大規模モデルトレーニングにアクセスできるようになることです。

LiBai モデルライブラリのアドレス: https://github.com/Oneflow-Inc/libai。

それで、具体的にはどうやってこれを実現するのでしょうか?以下では、上記の分散トレーニングツールをトレーニングのパフォーマンス、使いやすさなどの観点で比較し、次回分散トレーニングを行う際のツール選択の参考ガイドを提供します。

ワンクリック自動分散トレーニング、Megatron-LMやDeepSpeedを上回るパフォーマンス

シンプルで効率的な分散モデルトレーニングツールボックスとして、LiBai には次の 6 つの機能があります。

シングルカードコードから分散コードへのスムーズな拡張をサポートします。 LiBai の組み込みモデルは PyTorch と一貫性があり、学習と使用のコストを大幅に削減します。シンプルな構成で、あらゆる規模の並列処理に簡単に拡張できます。つまり、単一のカードに新しい機能を追加し、モデルをデバッグし、コードを実行してから、分散トレーニングにスムーズに移行できるということです。分散トレーニングをまったく構成したくない、または手動で構成した分散トレーニングが遅すぎると感じる場合は、分散ホスティング機能を試すことができます。自動並列パッケージ (https://libai.readthedocs.io/en/latest/tutorials/basics/Auto_Parallel.html) をインストールし、LiBai で線グラフ.auto_parallel=True を構成するだけです。分散を気にすることなく、モデル自体に集中して、より高速なトレーニング速度を得ることができます。
ハグフェイスに対応しています。 OneFlow と PyTorch は API レベルで高い互換性があります。簡単なコード変更で Hugging Face モデルをインポートできます。OneFlow を Torch としてインポートし、LiBai のデータ並列性、自動混合精度、Activation Checkpoint、ZeRO などのメカニズムに基づいて大規模モデルをトレーニングするだけです。モデルの個々のレイヤーを LiBai の組み込みレイヤーに置き換えると、3D 並列処理を使用して大規模なモデルをトレーニングできます。
モジュラー設計。 LiBai の実装では、モデル構築のために再利用可能な基本コンピューティングモジュールが提供されるだけでなく、データの読み込み、トレーニングロジック、指標の計算なども抽象化およびモジュール化されているため、ユーザーは自分のニーズに合わせて書き換え、トレーニング用のプラグインとして LiBai のトレーニングシステムに統合できます。
箱から出してすぐに使えます。大規模モデルのトレーニングには通常、いくつかのテクノロジへの依存が必要です。LiBai は、混合精度トレーニング、勾配再計算、勾配累積、ZeRO などの機能を提供しており、これらはデータ並列処理、モデル並列処理、パイプライン並列処理と簡単に組み合わせることができます。
すぐに実験を再現します。 OneFlow チームは、Detectron2 LazyConfig (https://github.com/facebookresearch/detectron2/blob/main/docs/tutorials/lazyconfigs.md) を参考にして LiBai の設定システムを構築しました。従来の argparse や yacs ベースの設定方法と比較すると、LiBai の設定システムはより柔軟で、Python 構文を使用して全体の構築を完了するため、新しいパラメータやモジュールを追加するのに非常に便利です。対応するモジュールをインポートするだけで、新しいモジュールの追加が完了します。同時に、トレーニング設定を yaml ファイルにシリアル化して保存することもできます。これにより、ファイル内で直接キーワードで設定項目を検索するのに便利です。ユーザーが以前の実験の結果を再現したい場合は、保存した config.yaml をトレーニング設定として直接渡すことができます。多くのスクリプトを含むファイルを保持すると、有効な変更を確認するのに役立たず、実験を再現するときに実験設定が混乱しやすくなります。
効率的なパフォーマンス。 Megatron-LMとの厳密なカーネルアライメントにより、さまざまなカーネル融合操作が実現されています。同時に、OneFlow静的グラフの設計により、LiBaiはシングルカードパフォーマンスとさまざまな並列処理の組み合わせの効率の点で、NVIDIAの高度に最適化されたMegatron-LMやMicrosoftのDeepSpeedよりも優れています。

OneFlow SBP がさまざまな並列テクノロジをネイティブにサポートしているため、LiBai はアルゴリズム記述と並列システムの分離を実現しました。わずか 30,000 行強のコードで、NVIDIA Megatron-LM と Microsoft DeepSpeed という 2 つの一般的なソリューションで合計 100,000 行のコードを必要とする機能を実現しました。

データは雄弁に語っています。以下の実験データはすべて、同じハードウェア環境、サードパーティの依存関係 (CUDA、cuDNN など)、パラメーター、ネットワーク構造に基づいています。さまざまなモデルでの LiBai と Megatron-LM のパフォーマンスを包括的に比較しています (すべてのパフォーマンス結果は公開されており、再現可能です、https://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html)。今後、OneFlow チームは、より大規模なクラスターでの LiBai のパフォーマンスを発表する予定です。

Megatron-LM 修正コミット: https://github.com/NVIDIA/Megatron-LM/commit/e156d2fea7fc5c98e645f7742eb86b643956d840 。
LiBai コミット: https://github.com/Oneflow-Inc/libai/commit/9fc504c457da4fd1e92d854c60b7271c89a55222 。
OneFlow コミット: https://github.com/Oneflow-Inc/oneflow/commit/55b822e4d3c88757d11077d7546981309125c73f 。

データ並列処理

注: 各パラメータグループの意味は以下のとおりです。
DP データ並列性、MP モデル並列性、PP パイプライン並列性、2D 並列性、3D 並列性。
fp16: 混合精度トレーニングを有効にする (amp) nl: レイヤー数 (パイプラインの並列サイズ = 8 の場合、各ステージで計算用の相対的なレイヤー数を使用できるように、レイヤー数を 24 から 48 に調整します)。
ac: アクティベーションチェックポイントを有効にする mb: GPU あたりのマイクロバッチサイズ gb: グローバルバッチサイズの合計 dxmxp。ここで、d = データ並列性 (データ並列サイズ)、m = モデル並列性 (テンソルモデル並列サイズ)、p = パイプライン並列性 (パイプラインモデル並列サイズ)。
1n1g は 1 台のマシンと 1 枚のカード、1n8g は 1 台のマシンと 8 枚のカード、2n8g は 2 台のマシンとそれぞれ 8 枚のカード、合計 16 枚のカード、4n8g は 4 台のマシンと合計 32 枚のカードを意味します。grad_acc_num_step = global_batch_size / (micro_batch_size * data_parallel_size) 表示される結果は Throughout です。

(注: このグループのレイヤー数 = 24、アンプオン、1n1g マイクロバッチサイズ = 24、他のグループのマイクロバッチサイズ = 16)

(注: このグループのレイヤー数 = 24、アンプオン、1n1g マイクロバッチサイズ = 6、他のグループのマイクロバッチサイズ = 4)

モデルの並列処理

(注: レイヤー数 = 24、amp が有効、アクティベーションチェックポイントが有効、マイクロバッチサイズ = 128、グローバルバッチサイズ = 1024、grad acc ステップ = 8)

(注: レイヤー数 = 24、アンプがオン)

パイプライン並列

(注: 最初の 2 つのグループのレイヤー数は 24、グラデーションアクセラレーションステップは 8、最後のグループのレイヤー数は 48、グラデーションアクセラレーションステップは 16 で、両方ともアンプがオンになっており、アクティベーションチェックポイントがオンになっています)

2Dデータ並列処理 + モデル並列処理

(注: このグループでは、レイヤー数 = 24、amp が有効、アクティベーションチェックポイントが有効、マイクロバッチサイズ = 128、grad acc ステップ = 8)

(注: このグループでは、num レイヤー = 24、amp が有効、アクティベーションチェックポイントが有効、マイクロバッチサイズ = 32、grad acc ステップ = 8)

データ + パイプラインの並列処理

(注: このグループでは、レイヤー数 = 24、amp が有効、アクティベーションチェックポイントが有効、マイクロバッチサイズ = 32、grad acc ステップ = 8)

3D 平行性

(注: このグループでは、num レイヤー = 24、amp が有効、アクティベーションチェックポイントが有効、grad acc ステップ = 8)

(注: このグループでは、num レイヤー = 24、amp が有効、アクティベーションチェックポイントが有効、grad acc ステップ = 8)

上記のパフォーマンス比較データから、厳密に調整された実験環境では、Bert モデルと GPT-2 モデルにおいて、LiBai のトレーニング速度があらゆる面で Megatron-LM を上回っていることがわかります。

他の人が持っているものを私も持っています: LiBai vs. 他のトレーニングプログラム

前述のように、大規模モデルのトレーニングの問題を解決するために、業界ではすでにHugging Face、DeepSpeed、Megatron-LM、FairSeqなどの人気のソリューションがあります。OneFlowがLiBaiモデルライブラリを開発する必要があるのでしょうか？

次に、上記のモデルライブラリの長所と短所を詳細に比較して、判断してみましょう。

HuggingFace : 包括的な SOTA Transformer モデルを提供し、事前トレーニング済みモデルを使用した微調整を容易にし、開発者が事前トレーニング済みモデルを使用できるようにするための強力なコミュニティとエコシステムを提供します。ただし、データの並列処理のみに対応しており、モデルが単一の GPU のメモリ容量を超える場合は適用できません。また、ゼロからのトレーニングの速度にも制限があります。

FairSeq : 主にシーケンスモデルを対象としています。NLP と CV の統合の傾向により、CV モデルのサポートが不足しています。

Megatron-LM : PyTorch をベースにデータ並列、モデル並列、パイプライン並列を実装しています。高性能で、超大規模モデルの真のトレーニングに使用できます。

しかし、多くのカスタマイズが施されており、分散トレーニングに慣れていないアルゴリズムエンジニアにとっては学習と使用の敷居が高すぎます。基本的には分散の専門家だけが再利用できます。さらに、Megatron-LM が提供するモデルは Hugging Face が提供するモデルに比べてはるかに少ないため、PyTorch を使用して大規模なモデルを再現したいエンジニアは、他の分散型の専門家が Megatron-LM に基づいて必要なモデルを実装するまで待たなければなりません。

DeepSpeed : PyTorch に基づくモデルメモリ最適化に関連するディープカスタマイズライブラリ。分散トレーニング、混合精度トレーニング、ZeRO などのテクノロジを提供し、メモリオーバーヘッドを効果的に節約し、データ並列処理で大規模モデルを効果的にトレーニングできるようにします。ただし、DeepSpeed はモデル並列処理をサポートしていません。モデルの一部の層のパラメータが大きすぎて単一の GPU のビデオメモリを超える場合、または DeepSpeed のセグメンテーション方法によって発生する通信効率が最適でない場合は、モデル並列処理 (テンソル並列処理、パイプライン並列処理) を使用する方がよいでしょう。現時点では、Megatron-LM のみを、侵入的な元のコードの変更と組み合わせて、ニーズを満たすことができます。

PyTorchエコシステムにおける大規模モデルトレーニングの創始者であるMegatron-LMとDeepSpeedに加えて、国内外の多くの著名な機関もFairSeqなどのいくつかの大規模モデルトレーニングライブラリを開発し、リリースしています。これらのライブラリの分散コア機能はすべてMegatron-LMとDeepSpeedに基づいていることを指摘しておく必要があります。

LiBai モデルライブラリがユニークなのは、上記の分散トレーニングツールの単純なアップグレードやパッケージ化ではなく、OneFlow の分散およびグラフコンパイラ機能に基づいて構築された大規模な事前トレーニング済みモデル開発キットである点です。このようにしてのみ、LiBai はパフォーマンスの面だけでなく、分散型の使いやすさの面でも他の追随を許さないものになります。

互換性。 PyTorch ベースで実装された現在の SOTA モデルと実質的に互換性があるため、ユーザーはモデルを迅速に移行できます。
効率。 1 枚のカードでも複数のカードでも、LiBai を使用することでユーザーはトレーニングの効率を向上させることができます。
使いやすさ。 LiBai は優れた拡張性を備えており、ニーズに応じてモデルを簡単に変更したり、新しい機能を追加したり、プロトタイプ機能の開発をより迅速に完了したりできます。ほとんど気づかないうちに、学習コストゼロで、ユーザーが分散型ディープラーニングトレーニングの敷居を大幅に下げるのに役立ちます。ユーザーが LiBai を使用して新しいモデルや新しい機能を開発する場合、単一の GPU のプログラミング方法を知っていれば、分散型トレーニング用にコードを書き直すことなく、大規模な GPU クラスターに自動的に拡張できるため、開発効率が向上します。

上記の比較を読んで、LiBai は AI エンジニアが分散トレーニングを行うための優れた選択肢になると思います。どう思いますか?

LiBaiは一般的な並列トレーニング戦略をすべてサポートしています

大規模モデルの分散トレーニングは、データ並列処理、モデル並列処理 (テンソル/モデル並列処理)、パイプライン並列処理などの複数の並列戦略を伴う複雑な問題です。LiBai モデルライブラリは、これら 3 つの一般的な並列戦略と、これらの並列戦略の任意の組み合わせをサポートしています (並列戦略の基本概念: https://docs.oneflow.org/master/parallelism/01_introduction.html)。

これらの並列戦略を自分で実装するのは、かなり面倒です。たとえば、自動混合精度トレーニングを使用するには、Apex の設定方法を学ぶ必要があります。データ読み込みパイプラインをサポートするには、DALI の設定方法を学ぶ必要があります。ZeRO を使用してビデオメモリの使用量を削減するには、DeepSpeed の設定方法を学ぶ必要があります... しかし、LiBai には複数の並列戦略が組み込まれており、スケーラビリティも優れているため、このような問題をまったく心配する必要はありません。

以下は、LiBai のさまざまな並列メソッドの例です。

ユニバーサル並列実装

OneFlow の SBP インターフェースを使用すると、ユーザーは自分のニーズと GPU のグループ化の配置に応じてネットワーク内の入力または重みを簡単に分割し、データまたはテンソルの並列処理を実現できます。

LiBai のレイヤーモジュール (libai.layers) には、一般的に使用される Linear、MLP、Transformer モジュールなど、さまざまな並列戦略に適応できる一連のネットワークレイヤーが組み込まれています。LiBai のレイヤーを使用して構築されたニューラルネットワークは、構成ファイルで分散構成のハイパーパラメータを調整するだけで、純粋なデータ並列処理、純粋なテンソル並列処理、およびデータとテンソルの混合並列トレーニング戦略を簡単に実装できます。

分散構成の形式は次のとおりです。

 # configs / common / train.py
 分散引数
辞書= 辞書(
 データパラレルサイズ= 1 、
 テンソルパラレルサイズ= 1 、
 パイプライン並列サイズ= 1 、
 ）

Data_parallel_size と tensor_parallel_size は、入力データとモデルの重みを異なる GPU グループに分割する方法を制御するために使用されます。ユーザーは、LiBai の組み込みレイヤーモジュールを使用してニューラルネットワークを構築した後、独自のトレーニング構成ファイルで分散ハイパーパラメータを変更して、さまざまな並列トレーニング戦略を実装できます。上の図では、すべての値が 1 に設定されており、単一のカードで実行されていることを示しています。ユーザーが 8 枚のカードのマシンを持っていると仮定して、この構成ファイルを変更してデータ並列処理、テンソル並列処理、パイプライン並列処理を実装する方法について説明します。

具体的な操作については、LiBai 分散構成ドキュメント (https://libai.readthedocs.io/en/latest/tutorials/basics/Distributed_Configuration.html) を参照してください。

純粋なデータ並列処理と純粋なモデル並列処理

ユーザーが 8 つの GPU で純粋なデータ (またはモデル) の並列トレーニングを実行する場合、トレーニング構成ファイル内の分散ハイパーパラメータを上書きするだけで済みます。

純粋なデータ並列処理

 # あなたのconfig.py
 libai.config からget_config をインポートします
train = get_config ( "common/train.py" ) 。train
 トレーニング. 分散. データ並列サイズ= 8

トレーニング中、同じモデルが異なるランクにコピーされ、各ランクが入力データの一部を処理してデータ並列トレーニングを実現します。

純粋なモデル並列処理

 # あなたのconfig.py
 libai.config からget_config をインポートします
train = get_config ( "common/train.py" ) 。train
 トレーニング. 分散. テンソル並列サイズ= 8

この場合、モデルは自動的に 8 つの GPU に分割され、各 GPU にはモデル全体の構造の一部のみが含まれるため、モデルの並列トレーニングが実現します。

データとモデルのハイブリッド並列トレーニング

ユーザーが 8 つの GPU で混合データとモデルの並列トレーニングを実行する場合、トレーニング構成ファイル内の分散ハイパーパラメータに次の簡単な変更を加えるだけです。

 # あなたのconfig.py
 libai.config からget_config をインポートします
train = get_config ( "common/train.py" ) 。train
 トレーニング. 分散. データ並列サイズ= 2
 トレーニング. 分散.tensor_parallel_size = 4

この場合、LiBai は GPU を自動的にグループ化します。8 つの GPU に [0、1、2、3、4、5、6、7] という番号を付けます。data_parallel_size=2 および tensor_parallel_size=4 が設定されている場合、8 つの GPU は実行時に自動的にグループ化され、[[0、1、2、3]、[4、5、6、7]] と表現されます。ここで、[0、1、2、3] はグループ、[4、5、6、7] はグループです。実行中、グループ間でデータ並列トレーニングが実行され、グループ内でモデル並列トレーニングが実行されます。

パイプライン並列構成

パイプライン並列処理の核となる概念は、ネットワークを複数のステージに分割し、異なるステージを異なる GPU に分散し、各ステージの計算結果を次のステージに渡して計算し、最終的にリレー方式でトレーニングを完了するという、簡単にまとめることができます。パイプラインの並列処理の詳細については、https://docs.oneflow.org/master/parallelism/01_introduction.html#_6 を参照してください。

シンプルなパイプライン並列構成

LiBai では、配置パラメータを設定することで、ネットワークの異なるレイヤーを異なる GPU に割り当てることができます。配置パラメータの値は、libai.utils.distributed の get_layer_placement() インターフェイスを通じて簡単に設定できます。LiBai は、設定ファイル (config) の分散構成に従ってステージを自動的に分割し、異なるステージに異なる配置を自動的に割り当てます。したがって、ネットワークの各レイヤーの配置を設定するだけで、分散構成と組み合わせることで、パイプライン並列構成を簡単に実装できます。

ほとんどのネットワークでは、線形層がネットワークのヘッドとして使用され、分類やその他のタスクのネットワークの最終結果を生成します。そこで、線形層を例に、LiBai で最も単純なパイプライン並列構成方法を簡単に紹介します。

 libai.layers からLinear をインポートします
self.head = Linear ( hidden_size , num_classes ) のようになります。

ネットワークモジュールの配置を構成する

LiBai では、ネットワークレイヤーを対応する配置に割り当てる方法が 2 つあります。

1. to_global インターフェイスと get_layer_placement() を組み合わせて配置を手動で指定します。ここでは、get_layer_placement(-1) を設定して、ヘッドレイヤーを最後のリレーの配置に設定します。

 libai.layers からLinear をインポートします
libai.utils.distributed をdist としてインポートします。
 self.head = Linear ( hidden_size , num_classes ) .to_global ( placement = dist.get_layer_placement ( -1 ) )

2. (推奨) libai.layers に実装されているモジュールには、layer_idx パラメータが付属しています。layer_idx パラメータを直接設定して、このレイヤーの配置を指定できます。

 libai.layers からLinear をインポートします
self.head = Linear ( hidden_size , num_classes , layer_idx = - 1 ) です。

入力データの配置を設定する

ネットワーク内のモジュールの配置を構成した後、入力データの配置も指定する必要があります。これは、入力とネットワークが同じステージにある場合にのみ計算を実行できるためです。最も直感的な方法は、入力とネットワークに同じ配置を構成することです。これは、to_global と get_layer_placement() を組み合わせることで実現できます。

 クラスMyModule ( nn . Module ):
 def __init__ ( self 、 ... * 、 layer_idx ):
 ...
 自己.layer_idx = レイヤー_idx
 ...
 def forward ( self 、 input_data ):
 input_data = input_data.to_global ( placement = dist.get_layer_placement ( self.layer_idx ) )
 ...

設定ファイルと組み合わせることで、単純なパイプライン並列処理を簡単に実現できます。

ネットワーク内のさまざまなレイヤーの配置と入力の配置を構成した後、パイプラインの並列処理を実行する前に、ユーザーは構成ファイル (config) を調整するだけで済みます。ネットワーク内のレイヤーの数は事前にわかっている必要があり、構成ファイルの pipeline_num_layers を調整する必要があります。

 # パイプラインのステージ数を2 に設定する
トレーニング. 分散. パイプライン並列サイズ= 2
 # パイプラインのモデルレイヤーを設定する
トレーニング. 分散. パイプライン_num_layers = 隠しレイヤー

1F1B は、PipeDream (https://arxiv.org/pdf/1806.03377.pdf) で導入された新しいパイプライン並列トレーニング方法で、ビデオメモリをより効率的に節約し、リソースを有効活用できます。 LiBai は、この 1F1B 戦略も比較的簡単にサポートできます (https://github.com/Oneflow-Inc/libai/blob/main/docs/source/tutorials/advanced_tutorials/customize_dataloader.md)。

3D並列処理の実装

データとモデルの混合並列処理とパイプライン並列処理を習得した後は、データ + モデル + パイプライン並列処理を構成するのは、上記のさまざまな並列処理の変更を組み合わせるだけです。

 # あなたのconfig.py
 libai.config からget_config をインポートします
train = get_config ( "common/train.py" ) 。train
 トレーニング. 分散. データ並列サイズ= 2
 トレーニング. 分散. テンソル並列サイズ= 2
 トレーニング. 分散. パイプライン並列サイズ= 2
 hidden_layers = 8 #ネットワーク内の層の数
トレーニング. 分散. パイプライン_num_layers = 隠しレイヤー

引き続き 8 つの GPU を例にとると、data_parallel_size、tensor_parallel_size、pipeline_parallel_size を 2 に設定すると、実行中にユーザーが設定した pinepine_num_layers に従って、モデルは GPU 上で自動的に分割されます。

上記の構成を例にとると、モデルは GPU [0、1、2、3] と [4、5、6、7] の 2 つのステージに分割されます。このうち、ステージ0はGPU [0, 2]と[1, 3]でデータ並列、GPU [0, 1]と[2, 3]でモデル並列、ステージ1はGPU [4, 6]と[5, 7]でデータ並列、GPU [4, 5]と[6, 7]でモデル並列となります。

カスタム並列トレーニング

上記の紹介によると、LiBai はユーザーが呼び出すために libai/layers/ の下にパッケージ化されたモジュールを提供します。これらのモジュールを組み合わせることで、ユーザーは独自の並列ネットワークを構築できます。

LiBai のモジュールがユーザーのニーズを満たせない場合、ユーザーは並列戦略を非常に便利にカスタマイズすることもできます。 scatter -> forward -> Reduce などの一連の複雑な通信操作を手動で挿入する必要がある PyTorch とは異なり、LiBai では、ユーザーはテンソルを初期化するときに sbp と配置を定義するだけで、スタンドアロンコードを記述するのと同じように独自の並列コードを実行できます。 (sbp と配置の詳細については、https://docs.oneflow.org/master/parallelism/04_2d-sbp.html を参照してください)。

たとえば、ユーザーが 4 カードトレーニングを実行すると、ネットワークの中間結果は (16, 8) の形状を持つ 2D 並列テンソルを持ち、LiBai の次の図に示すように GPU 上で分割されます。このテンソルの配置分布はranks=[[0, 1],[2, 3]]であり、SBPは(S[0], S[1])または(S[1], S[0])です。

 [ |
 X00 GPU0 | X01 GPU1
 --------------------------
 X10 GPU2 | X11 GPU3
 | ]

このうち、Xij の形状は (8, 4) であり、各カードに均等に分布しています。このテンソルにランダムなノイズを追加したい場合は、LiBai で次のコードを簡単に追加できます。

LiBai は、1D 並列要件と互換性を持つように dist.get_nd_sbp() をカプセル化し、dist.get_layer_placement() はパイプライン並列の構成を容易にします。ほとんどの場合、ユーザーは次のコードを直接参照できます。

 # テスト.py
 oneflowをフローとしてインポートする
omegaconf からDictConfigをインポートする
oneflow からインポートnn
 libai.utils からインポートし、 dist  として配布します
cfg = 辞書設定(
 dict ( データパラレルサイズ= 2 、 テンソルパラレルサイズ= 2 、 パイプラインパラレルサイズ= 1 ))
 dist . setup_dist_util ( cfg )
 クラスNoise ( nn . モジュール):
 def __init__ ( 自己):
 スーパー()。__ init__ ()
 self.noise_tensor = flow.randn ( 
 16 , 8 ,
 sbp = dist.get_nd_sbp ( [ flow.sbp.split ( 0 ) ,flow.sbp.split ( 1 ) ]) ,
 配置= dist 。 get_layer_placement ( layer_idx = 0 )
 ）
 # 次のように変更することもできます
# self.noise_tensor = flow.randn ( 
 # 16 , 8 ,
 # sbp = ( フロー.sbp.split ( 0 ), フロー.sbp.split ( 1 ) ),
 # placement = flow.placement ( "cuda" 、 ranks = [ [ 0 , 1 ],[ 2 , 3 ]])
 # )
 def forward ( self , x ):
 x + self.noise_tensor を返す
ノイズ= ノイズ()
 x = フローゼロ(
 16 , 8 ,
 sbp = ( フロー. sbp . スプリット( 0 )、 フロー. sbp . スプリット( 1 ))、
 配置= フロー. 配置( "cuda" 、 ランク= [[ 0 、 1 ]、[ 2 、 3 ]])
 ）
 y = ノイズ( x )
 print ( f "ランク: {flow.env.get_rank()}、グローバルテンソル: 形状 {y.shape} sbp {y.sbp} 配置 {y.placement}、ローカルテンソル形状: {y.to_local().shape}" )

コマンドを実行:

 python3 -m oneflow.distributed.launch --nproc_per_node 4 テスト.py

次の出力は、形状に応じた各ランクのテンソルの分布と、全体的な観点からのテンソルの情報を示しています。

 ランク: 2 、 グローバルテンソル: 形状oneflow.Size ([ 16 , 8 ] ) sbp ( oneflow.sbp.split ( axis = 0 )、 oneflow.sbp.split ( axis = 1 ) ) 配置oneflow.placement ( type = "cuda" 、 ranks = [[ 0 , 1 ] 、 [ 2 , 3 ]])、 ローカルテンソルの形状: oneflow.Size ( [ 8 , 4 ])
 ランク: 3 、 グローバルテンソル: 形状oneflow.Size ([ 16 , 8 ] ) sbp ( oneflow.sbp.split ( axis = 0 )、 oneflow.sbp.split ( axis = 1 ) ) 配置oneflow.placement ( type = "cuda" 、 ranks = [[ 0 , 1 ] 、 [ 2 , 3 ]])、 ローカルテンソルの形状: oneflow.Size ( [ 8 , 4 ])
 ランク: 1 、 グローバルテンソル: 形状oneflow.Size ([ 16 , 8 ] ) sbp ( oneflow.sbp.split ( axis = 0 )、 oneflow.sbp.split ( axis = 1 ) ) 配置oneflow.placement ( type = "cuda" 、 ranks = [[ 0 , 1 ]、[ 2 , 3 ]])、 ローカルテンソルの形状: oneflow.Size ( [ 8 , 4 ] )
 ランク: 0 、 グローバルテンソル: 形状oneflow.Size ([ 16 , 8 ] ) sbp ( oneflow.sbp.split ( axis = 0 ) 、 oneflow.sbp.split ( axis = 1 ) ) 配置oneflow.placement ( type = "cuda" 、 ranks = [ [ 0 , 1 ]、[ 2 , 3 ]])、 ローカルテンソル形状: oneflow.Size ( [ 8 , 4 ])