OpenAI: 大規模ニューラルネットワークをトレーニングするための 4 つの基本手法

この記事はLeiphone.comから転載したものです。転載する場合は、Leiphone.com公式サイトにアクセスして許可を申請してください。

大規模ニューラル ネットワークは、現在の人工知能分野で注目されているトピックの 1 つです。では、大規模モデルをどのようにトレーニングするのでしょうか?

最近、大規模な事前トレーニング済みモデル GPT-3 を発表した OpenAI は、GPU をベースにしたメモリを節約する 4 つの並列トレーニング手法を紹介するブログ記事を公開しました。

• データ並列処理 - 同じバッチの異なるサブセットを異なる GPU で実行します。
• パイプラインの並列処理 - モデルの異なるレイヤーを異なる GPU で実行します。
• テンソル並列処理 - 行列乗算を GPU に分割するなど、数学演算を個別の演算に分解します。
• Mixture of Experts (MOE) - 各レイヤーのごく一部だけを通じて各例を処理します。

図のキャプション: 3 層モデルにおけるさまざまな並列戦略。各色は層を表し、点線は異なる GPU を区別します。

1.データの並列処理

「データ並列トレーニング」とは、同じパラメータを複数の GPU (「ワーカー」と呼ばれることが多い) にコピーし、各 GPU に異なるサンプルセットを割り当てて同時に処理することを意味します。

データの並列処理だけでも、モデルを単一の GPU のメモリに収める必要がありますが、計算に複数の GPU を活用する場合、パラメーターの複数のコピーを保存するコストがかかります。ただし、そうは言っても、使用の合間にパラメータを一時的に CPU メモリにオフロードするなど、GPU で使用できる有効な RAM を増やす戦略はあります。

各データ並列ワーカーがパラメータのコピーを更新する際、各ワーカーが引き続き同様のパラメータを持つように、ワーカー同士が調整する必要があります。最も単純なアプローチは、ワーカー間の「ブロッキング通信」を導入することです。

ステップ 1: 各ワーカーの勾配を個別に計算します。

ステップ 2: 異なるワーカーの勾配を平均します。

ステップ 3: 各ワーカーで同じ新しいパラメータを個別に計算します。

ステップ 2 は、大量のデータ (ワーカー数とパラメータ サイズを掛け合わせた値に比例) を転送する必要があるブロッキング平均であり、トレーニングのスループットが低下する可能性があります。この損失を排除できるさまざまな非同期同期方式がありますが、学習効率が犠牲になるため、実際には同期方式が一般的に使用されます。

2つのパイプラインを並列に

パイプライン並列トレーニングでは、研究者はモデルの連続ブロックを GPU に分割し、各 GPU はパラメータのごく一部のみを保存するため、同じモデルに対して各 GPU が消費するメモリが比例して削減されます。

大規模なモデルを連続するレイヤーのチャンクに分割するのは簡単ですが、レイヤーの入力と出力の間には連続的な依存関係があるため、単純な実行では、ワーカーが前のマシンの出力が入力として使用されるのを待機する間に、かなりのアイドル時間が発生する可能性があります。これらの待機時間のブロックは「バブル」と呼ばれ、アイドル状態のマシンで実行できたはずの無駄な計算です。

図のキャプション: モデルが垂直に 4 つのパーティションに分割された、単純なパイプライン並列設定の図。ワーカー 1 は最初のレイヤー (入力に最も近い) のモデル パラメータをホストし、ワーカー 4 は 4 番目のレイヤー (出力に最も近い) をホストします。 「F」、「B」、「U」はそれぞれ前方、後方、更新操作を表します。下付き文字は、どのワーカーで操作を実行するかを示します。順次的な依存関係により、データは一度に 1 つのワーカーによって処理され、大量のアイドル時間の「バブル」が発生します。

データ並列処理のアイデアを再利用して、各ワーカーが一度にデータ要素のサブセットのみを処理できるようにすることで、タイムバブルの作成コストを削減し、新しい計算を待機時間と巧みに重ねることができます。基本的な考え方は、バッチを複数のマイクロバッチに分割することで、各マイクロバッチの処理速度が比例して速くなり、次のマイクロバッチが利用可能になるとすぐに各ワーカーが作業を開始し、パイプラインの実行を高速化するというものです。十分なマイクロバッチがあれば、ステップの開始時と終了時の「バブル」を最小限に抑えながら、ほとんどの時間ワーカーを活用できます。勾配はミニバッチ全体で平均化され、すべてのミニバッチが完了した後にのみパラメータが更新されます。

モデルが分割されるワーカーの数は、多くの場合「パイプラインの深さ」と呼ばれます。

フォワード パス中、ワーカーはレイヤー ブロックの出力 (「アクティベーション」と呼ばれる) を次のワーカーに送信するだけです。バックワード パス中、ワーカーはそれらのアクティベーションの勾配のみを前のワーカーに送信します。これらのチャネルをどのように配置し、ミニバッチ間でグラデーションを集約するかについては、大きな設計余地があります。たとえば、GPipe アプローチでは、各ワーカー プロセスが順方向パスと逆方向パスを連続的に実行し、最後に複数のマイクロバッチからの勾配を同期的に集約します。一方、PipeDream では、各ワーカーが順方向パスと逆方向パスを交互に処理するように調整します。

図のキャプション: バッチごとに 4 つのマイクロバッチを使用した GPipe と PipeDream パイプライン ソリューションの比較。ミニバッチ 1 ～ 8 は、連続する 2 つのデータ バッチに対応します。図中の「番号」はどのマイクロバッチが操作されているかを示し、下付き文字はワーカー ID を示します。 PipeDream は、古いパラメータを使用して一部の計算を実行することで、より高い効率を実現することに注意してください。

3テンソル並列処理

パイプラインの並列処理では、モデルをレイヤーごとに「垂直に」分割し、レイヤー内で特定の操作を「水平に」分割することもできます。これは、テンソル トレーニングと呼ばれることがよくあります。

多くの最新モデル (Transformer など) では、計算上のボトルネックは、アクティベーション バッチ マトリックスと大きな重みマトリックスを乗算することです。行列の乗算は、行と列のペア間のドット積と考えることができます。異なる GPU で個別のドット積を計算したり、異なる GPU で各ドット積の一部を計算して結果を合計したりすることも可能です。戦略に関係なく、重みマトリックスを均等なサイズの「部分」に分割し、各部分を異なる GPU でホストし、その部分を使用して完全なマトリックス積の関連部分を計算してから、結果を結合するために通信することができます。

一例として、Transformer の自己注意層と MLP 層内で行列乗算を並列化する Megatron-LM があります。 PTD-P はテンソル、データ、パイプラインの並列処理を使用し、パイプライン スケジューリングによって各デバイスに複数の個別のレイヤーを割り当てて、ネットワーク通信の増加を犠牲にしてバブル損失を削減します。

場合によっては、ネットワーク入力を次元間で並列化することができ、その結果、相互通信に比べて高度な並列計算が可能になります。シーケンス並列処理とは、入力シーケンスを時間的に複数のサブ例に分割し、より細かい例で計算を続行できるようにすることで、ピーク時のメモリ消費を比例して削減するという考え方です。

4専門家の組み合わせ (文部科学省)

専門家の混合 (MoE) アプローチを使用すると、ネットワークのごく一部のみが、1 つの入力に対する出力の計算に使用されます。

一例として、複数の重みセットを用意し、ネットワークが推論時にゲーティング メカニズムを通じて使用する重みセットを選択できるようにするという方法があります。これにより、計算コストを増やすことなく、より多くのパラメーターを使用できるようになります。重みの各セットは「エキスパート」と呼ばれ、ネットワークが各エキスパートに専門的な計算とスキルを割り当てることを学習することが期待されています。異なるエキスパートが異なる GPU をホストできるため、モデルに使用される GPU の数を明確にスケールアップできます。

キャプション: ゲーティング ネットワークは n 人の専門家のうち 2 人だけを選択しました。

GShard は、MoE Transformer のパラメータを 6000 億パラメータに拡張し、MoE レイヤーのみが複数の TPU デバイスに分割され、他のレイヤーは完全に複製されます。 Switch Transformer は、1 つの入力を 1 つのエキスパートにルーティングすることで、モデル サイズをより高いスパース性を持つ数兆個のパラメータに拡張します。

5その他のメモリ節約設計

より大規模になるニューラル ネットワークのトレーニングをより扱いやすくする計算戦略は他にも多数あります。例えば：

勾配を計算するには、元のアクティベーションを保存する必要があり、デバイスの RAM を大量に消費します。チェックポイント (アクティベーション再計算とも呼ばれる) は、アクティベーションのサブセットを保存し、バックワード パス中に中間アクティベーションをジャスト イン タイムで再計算します。これにより、最大で 1 つの完全なフォワード パスの追加の計算コストで大量のメモリを節約できます。また、選択的アクティベーション再計算によって計算コストとメモリ コストを継続的にトレードオフすることもできます。選択的アクティベーション再計算では、アクティベーションのサブセットをチェックしますが、保存コストは比較的高くなりますが、計算コストは​​低くなります。

混合精度トレーニングは、より低い精度の数値 (最も一般的には FP16) を使用してモデルをトレーニングするプロセスです。最新のアクセラレータは、より低い精度の数値を使用してより高い FLOP カウントを達成し、デバイスの RAM も節約できます。適切な注意を払えば、精度をほとんど損なうことなくモデルを作成できます。

オフロードとは、未使用のデータを CPU または異なるデバイス間で一時的にオフロードし、必要に応じて読み戻すことです。単純な実装ではトレーニングの速度が大幅に低下しますが、高度な実装ではデータを事前に取得するため、デバイスが待機する必要がなくなります。このアイデアの 1 つの実装が ZeRO です。これは、パラメーター、勾配、およびオプティマイザーの状態をすべての利用可能なハードウェアに分割し、必要に応じて実現します。

メモリ効率の高いオプティマイザ メモリ効率の高いオプティマイザは、Adafactor などのオプティマイザによって維持される実行時状態のメモリ フットプリントを削減するために提案されています。

圧縮は、ネットワーク内の中間結果を保存する場合にも使用できます。たとえば、Gist は後方パス用に保存されたアクティベーションを圧縮し、DALL-E は同期する前に勾配を圧縮します。