序文

この記事は、「1000 億パラメータのオープンソース大規模モデル BLOOM の背後にある技術、これは英語の原文です」および関連論文の研究ノートとみなすことができますが、いくつかの詳細と誤りが修正され、多数の説明が修正されました。読みやすくするために追加しました明確で理解しやすい

BLOOMの前編とその背後にあるメガトロン・ディープスピード

1.1 BLOOM トレーニングの詳細: ハードウェア/チェックポイント/データセット

BLOOM のモデルアーキテクチャは GPT3 に非常に似ていますが、いくつかの改良が加えられています。176B BLOOM モデルのトレーニングは、2022 年 3 月から 7 月までに完了するまでに約 3.5 か月 (約 100 万計算時間) かかります。以下はそのトレーニングです。いくつかの詳細

トレーニングハードウェア

GPU: 384 個の NVIDIA A100 80GB GPU (48 ノード) + 32 個の予備 GPU
ノードあたり 8 つの GPU、4 つの NVLink カード間インターコネクト、4 つの OmniPath リンク
CPU: AMD EPYC 7543 32コアプロセッサ
CPU メモリ: ノードあたり 512GB
GPU メモリ: ノードあたり 640GB
ノード間の接続: Omni-Path Architecture (OPA) ネットワークカードが使用され、ネットワークトポロジはノンブロッキングファットツリーです。
NCCL - 通信ネットワーク: 完全専用のサブネットワーク
ディスク IO ネットワーク: 他のノードおよびユーザーと共有される GPFS

チェックポイント

主なチェックポイント
各チェックポイントには、fp32 の精度のオプティマイザ状態と bf16+fp32 の精度の重みが含まれており、2.3 TB のストレージ領域を占有します。bf16の重量だけを節約した場合、占有されるストレージ容量は329GBだけです。

データセット

41.5TB の高度に重複排除およびクリーンアップされた 46 言語のテキスト (350B トークンに変換)
モデルの語彙には 250,680 個のトークンが含まれています
詳細については、 BigScience Corpus A 1.6TB 複合多言語データセットを参照してください。

1.2メガトロン-ディープスピード：

176B BLOOM モデルは、次の 2 つの主要な手法を組み合わせたMegatron-DeepSpeedを使用してトレーニングされます。

Megatron-LM は、NVIDIA の応用深層学習研究チームによって開発された大規模で強力なトランスフォーマーモデルフレームワークです。対応する論文は「Megatron-LM: Training Multi-Billion Parameter Language Models using Model Parallelism」、これは Li Mu のビデオ解釈です、これはテキスト/コードの解釈です
DeepSpeed は、Microsoft によって開発されたディープラーニング最適化ライブラリであり、分散トレーニングをシンプル、効率的、効果的に行うことができます。

DeepSpeed チームは、以下の最初の項目と最後の 3 つの項目を組み合わせました。

Megatron-LM のテンソル並列処理 (モデル並列処理の一種として理解できるテンソル並列処理) では、
各テンソルが複数のブロックに分割されるため、テンソル全体を常駐させるのではなく、テンソルの各スライスが指定された GPU に配置されます。単一の GPU 上で。処理中、各シャードは別の GPU で個別に並列処理され、結果はステップの最後に同期されます。これは水平平行度と呼ばれます。
Zero Redundancy Optimizer (Microsoft DeepSpeed ライブラリのコアである略称 ZeRO)
も TP と同様のテンソルシャーディングを実行しますが、テンソル全体は順方向または逆方向の計算に間に合うように再構築されるため、モデルは存在しません変更する必要があります。また、限られた GPU メモリを補うためのさまざまなオフロード手法もサポートしています。
データ並列処理 (データ並列処理)
同じ設定とモデルが複数のコピーにレプリケートされ、各コピーには毎回異なるデータコピーが供給されます。処理は並行して実行され、各トレーニングステップの終了時にすべての共有が同期されます。
パイプライン並列処理 (パイプライン並列処理とも呼ばれる)
モデルは、複数の GPU に垂直に (つまり、レイヤーごとに) 分割され、1 つ以上のモデルレイヤーのみが 1 つの GPU に配置されます。各 GPU はパイプラインのさまざまなステージを並行して処理し、バッチの一部を処理します。

Megatron-Deepspeed という 3D 並列実装を開発しました。これにより、BLOOM など、1,000 億を超えるパラメータを持つ大規模な言語モデルの分散トレーニングがより簡単、より効率的、効果的になります。

BLOOM チームの BigScience バージョンのMegatron-DeepSpeed は、オリジナルの Megatron-DeepSpeed コードベースに基づいていますが、その上にかなりの数のコードが追加されていることに注意してください。

次の表は、BLOOM のトレーニングに使用した 2 つのフレームワークそれぞれのコンポーネントを示しています。

コンポーネント	ディープスピード	メガトロンLM
ゼロデータパラレル	はい
テンソル並列		はい
パイプラインパラレル	はい
BF16オプティマイザー	はい
CUDAフュージョンカーネル機能		はい
データローダー		はい

Megatron-LM と DeepSpeed の両方にはパイプライン並列処理と BF16 オプティマイザーの実装がありますが、ZeRO に統合されているため、DeepSpeed の実装を使用することに注意してください。

2 番目の部分は Tensor Parallelism (モデル並列処理の一種である Tensor Parallelism) です。

Tensor Parallelism (TP) では、各 GPU はテンソルの一部のみを処理し、特定の演算子が完全なテンソルを必要とする場合にのみ集約操作がトリガーされます。

ご存知のとおり、Transformer には 2 つの主要なモジュールがあります:一个自注意力层 + 残差连接，完全接続層 MLP + 残留接続

2019 年に、NVIDIA は Megatron-LM 論文 ( GPU クラスターでの効率的な大規模言語モデルのトレーニング) を通過させました。その内積部分は次のように書くことができます $Y = GeLU(XA)$ 。ここで、と $バツ$ は $Y$ 入力ベクトルと出力ベクトル、 $あ$ 重み行列です。

次の図 (図 1と表記)に示すように、行列の形式で表すと、行列の乗算が複数の GPU にどのように分割されるかが簡単にわかります。

2.1 MLP の並列化: 重み A 行列を垂直に切り、B 行列を水平に切り、最後に MERGE

上の概略図を過小評価しないでください。実際、そこには何度も熟考する価値のある多くの詳細が含まれています。詳細を次の図に示します (図 2 と表記)

入力の場合 $バツ$ 、その行数はバッチ $b$ サイズ $私$ 、列数は隠れ層の幅です。つまり、 $k$
その隠れ層のモジュールは実際には完全に接続された 2 つの層です。
最初の隠れ層の重みが $あ$ (行数が K、列数が K'、K' は通常 K の 4 倍) であると仮定すると、最初に行列乗算を実行し $X\cdot A$ 、次に $\シグマ$ GELU などの活性化関数を接続します。 ( GELU は ReLU の変曲点が滑らかに下降するのと似ています)
2 番目の隠れ層の重みが $B$ (行数 K'、列数 K ) であると仮定すると、最終的な $\sigma (X\cdot A)B = Y$
次に、分割の方法を見てみましょう。並列処理には複数の GPU を使用する方が良いでしょう。
入力データが比較的大きい場合は、最初にデータ並列処理を行うことを選択します。つまり、入力を分割します。モデル自体が比較的大きい $バツ$ 場合は、
モデルは並列、つまり行列 $あ$ が分割されており、分割方法は 2 つの
$\rightarrow$ タイプに分けられます。1 つ目 (上記の図 1の下部に相当します): 行列 $あ$ は行ごとに水平方向に分割 ( $バツ$ その後、対応する列が垂直方向に分割され、その結果、2 つの GPU 間の通信が必要になります)
$\rightarrow$ 2 番目のタイプ (上の図 1の上部に対応): 行列は $あ$ 次のように垂直方向に分解されます。上の図2に示すように、一方の列は青、もう一方の列は緑です(対応して行ごと $バツ$ に水平に分割するか、 $バツ$ 両方の GPU にコピーが存在する必要があり、現時点では追加の通信は必要ありません)。
2 番目の分割方法 (つまり、行列が $あ$ 列ごとに垂直に分割される) を確認した後、乗算して $バツ$ 大きな行列を取得し $X\cdot A$ 、 $B$ 行列を行ごとに切り取ります (上の図 2に示すように、2 番目の行列は行列です) B、青いものは GPU 0 に配置され、緑のものは GPU 1 に配置されます)、最後の行列の $シャー$ 列と行列の行 $B$ は行列の乗算を行い、取得されたサイズは $Y$ のサイズと一致しますが、結果は次のとおりです。 $1/N$ ( $N$ GPU の数を参照)
つまり、行列乗算を実行して最初に取得し $XA_1$ 、 $XA_n$ 次に出力ベクトルを取得することで、それらを個別に GeLU に入力し、最後に上記のステップ 5 の結果をマージして完全なベクトルを取得できます。上記の操作により、シーケンスの後に GPU を同期するたびに、任意の深さの MLP を更新できます。Megatron-LM 論文の著者は、これについて素晴らしい図解を提供しました (図 3 として示されています)。 $N$ $Y_1 Y_2 \cdots Y_n$ ：
$Y$

拆列-拆行
ここでは $f$ 、前方パスの恒等演算子と後方パスのすべてのリダクション $g$ ですが前方パスのすべてのリダクションと後方パスの恒等演算子です。

2.2 マルチヘッドアテンションレイヤーの並列化: 各ヘッドが個別に計算

マルチヘッドアテンションレイヤーは複数の独立したヘッドにより本質的に並列であるため、マルチヘッドアテンションレイヤーの並列化はさらに簡単になります。以下の図に示すように (図 4 とマークされています)

入力 $バツ$ 行列の場合、行数は依然としてバッチ $b$ サイズ $私$ (バッチサイズが 1 であると仮定)、列数はセルフ $k$ アテンションメカニズムのです (セルフアテンションメカニズムを忘れた場合は、詳細については、この記事を参照してください。Transformer Notesの 3 番目の部分)、入力は $バツ$ 以下に対応する 3 つのコピーにコピーされます: $バツ$ ベクトル $Q$ $K$ $V$ 行列 (3 つのクローンと同様)
マルチヘッドアテンションに関しては、ヘッドの次元はであると $k/h$ 仮定すると、その後、 $h=2$ 各ヘッドの入力行列内の各 $バツ$ 単語ベクトルに対して $Q$ 、それぞれのコンテキストのベクトルとスケーリングドット積を実行し、 $K$ 次にソフトマックスを実行して次のようになります。注意スコアまたは重みを取得し、それを $V$ 重み付き合計で実行して $l \times k/h$ 出力を取得し、
最後に $k\回k$ 射影を乗算して結果を取得します $l \times k$ 。
2 番目のヘッドの計算プロセスも同様です
。各ヘッドの計算は互いに影響を与えることなく独立して並列していることがわかります。これは、1 つのヘッドを GPU 0 (青で示されている) に配置でき、もう 1 つのヘッドが GPU 0 に配置できることを意味します。 GPU 1 (緑色で示されている) に配置され、各ヘッドの結果が最終的にすべて縮小されます。

全体のプロセスは、下の図に示されています (図 5 として示されています)。

以下については特別な考慮が必要です。

順方向伝播と逆方向伝播では層ごとに 2 つの all-reduce があるため、TP ではデバイス間の非常に高速な相互接続が必要です。したがって、非常に高速なネットワークがない限り、複数のノード間で TP を実行することはお勧めできません。BLOOM をトレーニングするためのハードウェア構成では、ノード間の速度は PCIe よりもはるかに遅くなります。実際、ノードに 4 つの GPU がある場合、最大 TP 次数は 4 の方が優れています。TP レベル 8 が必要な場合は、少なくとも 8 つの GPU を備えたノードを使用する必要があります。

このコンポーネントは Megatron-LM によって実装されます。Megatron-LM は最近、テンソル並列機能を拡張し、前述のセグメンテーションアルゴリズムを使用するのが難しい演算子のためにシーケンス並列機能を追加しました。たとえば、論文「LayerNorm Reducing Activation Recomputation in Large Transformer Models」には、このテクノロジの詳細情報が記載されています。シーケンスの並列処理は BLOOM のトレーニング後に開発されたため、BLOOM のトレーニングではこの手法は使用されません

2.3 入出力の並列化

次に、以下の図 (図 6 と表記) に示すように、入力と出力の並列化を見てみましょう。

入力の場合、 $バツ$ これ $b$ は $私$ バッチサイズとシーケンスの長さを乗じた行列で、文を 1 行ずつ格納します。埋め込み層は vocab (語彙に相当) の行と K の列を持つ語彙です。辞書全体を水平にカットすることができます
(たとえば、上半分を青いマークで GPU 0 に配置し、下半分を緑のマークで GPU 1 に配置)、テーブルを参照して出力を取得し $b\times l \times k$ ます
出力の行数は K で、 $b\×l$ 列数は K です。ボキャブラリを通過した後、 $b\times l \times v$ 出力が取得され、その左半分は GPU 0 に配置でき、右半分は GPU 1 に配置できます。出力の各ラインは水平方向に加算できますが、V は数万のサイズなど比較的大きい場合があります。もちろん、各 GPU は独自の部分をカウントでき
ます。 $b\times l \times v$

パート III データパラレルと ZeRO

数個の GPU しか持たないほとんどのユーザーは、おそらく DistributedDataParallel(DDP) に精通しているでしょう。これは、対応する PyTorch ドキュメントです。このアプローチでは、モデルが各 GPU に完全に複製され、各反復後にすべてのモデルが状態を相互に同期します。この方法では、より多くの GPU リソースを投資することでトレーニングを高速化し、問題を解決できます。ただし、モデルが単一の GPU に適合する場合にのみ機能するという制限があります。

3.1 ZeRO データ並列処理

3.1.1 ゼロ 1

2020年、Microsoft DeepSpeedチームは論文「ZeRO: Memory Optimizations Toward Training Trillion Parameter Models」を通じてZero Redundancy Optimizer (略してZeRO)を提案しましたが、最適化は永遠のテーマであるため、DeepSpeedチームは3つの論文を発表しました。ここ数年のZeRO関連の論文（最新記事で提案した手法をZeRO 3と呼びます）では、1つの目標を掲げて冗長なパラメータの削除、CPUやメモリの導入、NVMeの導入などの手法が提案されています。 Beginning to End: ビデオメモリの最適化を最後まで実行します。

以下の図 7 は、ZeRO データ並列処理をわかりやすく説明しています (このブログ投稿から)

比較的縦長なので理解に集中しにくいように思えますが、実は概念は非常にシンプルです。これは、各 GPU が完全なモデルパラメーター、勾配、およびオプティマイザーの状態を複製する代わりに、その一部のみを保存する点を除いて、単なる通常の DDP です。後続の実行中に、特定のレイヤーの完全なレイヤーパラメーターが必要になると、すべての GPU が同期して、不足している部分を互いに提供します。それ以上のことはありません。

3.1.2 ゼロ2

// 更新される

3.1.3 ゼロ3

更新される予定です。

以下は、8.25 午後 4 時に変更されます。

パイプラインの 4 番目の部分は並列です

単純なパイプライン並列処理 (単純な PP) は、モデルレイヤーを複数の GPU にグループ化して分散し、1 つの大きな複合 GPU であるかのように GPU から GPU にデータを単純に移動します。メカニズムは比較的単純です。目的のレイヤー .to() メソッドを対応するデバイスにバインドします。これで、データがこれらのレイヤーに出入りするたびに、レイヤーはデータをレイヤーと同じデバイスに切り替え、残りは同じままになります。

これは実際には垂直モデルの並列処理です。ほとんどのモデルのトポロジをどのように描画したかを覚えていると思いますが、実際にはモデルのレイヤーを垂直に分割しているからです。たとえば、以下の画像が 8 層モデルを示しているとします。

===================  ===================
|  0 | 1 | 2 | 3  |  |  4 | 5 | 6 | 7  |
===================  ===================
        GPU0                 GPU1

これを垂直に 2 つの部分に切り、レイヤー 0 ～ 3 を GPU0 に、レイヤー 4 ～ 7 を GPU1 に配置します。

ここで、データがレイヤー 0 からレイヤー 1 に、レイヤー 1 からレイヤー 2 に、レイヤー 2 からレイヤー 3 に渡されるとき、それは単一の GPU での通常の順方向パスとまったく同じです。ただし、データをレイヤー 3 からレイヤー 4 に渡す必要がある場合は、GPU0 から GPU1 に転送する必要があるため、通信オーバーヘッドが発生します。参加している GPU が同じ計算ノード (同じ物理マシンなど) 上にある場合、転送は非常に高速ですが、GPU が異なる計算ノード (複数のマシンなど) 上にある場合、通信オーバーヘッドが大幅に大きくなる可能性があります。

次に、レイヤー 4、5、6、7 は再び通常のモデルと同様になります。レイヤー 7 が完了したら、通常、ラベルがあるレイヤー 0 にデータを送り返す (またはラベルを最後のレイヤーに送信する) 必要があります。これで、損失を計算し、オプティマイザを使用してパラメータを更新できるようになります。

質問：

この方法が単純なパイプライン並列処理と呼ばれるのはなぜですか?また、その欠点は何ですか? 主な理由は、このスキームでは常に 1 つを除くすべての GPU がアイドル状態になるためです。したがって、4 つの GPU を使用すると、1 つの GPU のメモリ量がほぼ 4 倍になり、他のリソース (コンピューティングなど) はほとんど役に立たなくなります。デバイス間でデータをコピーするオーバーヘッドを追加します。したがって、単純なパイプラインを使用して 4 枚の 6GB カードを並列に使用すると、1 枚の 24GB カードと同じサイズのモデルを保持でき、データ転送のオーバーヘッドがないためトレーニングが高速になります。ただし、たとえば、40 GB カードを持っているが、45 GB モデルを実行する必要がある場合は、40 GB カードを 4 枚使用できます (ビデオメモリを必要とするグラデーションやオプティマイザ状態もあるため、これで十分です)。
エンベディングを共有するには、GPU 間でのコピーが必要になる場合があります。私たちが使用するパイプライン並列処理 (PP) は、上記の単純な PP とほぼ同じですが、受信バッチをマイクロバッチに分割し、異なる GPU が同時に計算プロセスに参加できるパイプラインを人為的に作成することで、GPU のアイドリング問題を解決します。

以下の図は GPipe 論文からのもので、上の部分は単純な PP スキームを表し、下の部分は PP メソッドを表します。

図の下半分から、PP のデッドゾーン (GPU がアイドル状態であることを意味します) が少ない、つまり「バブル」が少ないことが簡単にわかります。

図の 2 つのスキームの並列度は 4 です。つまり、パイプラインは 4 つの GPU で構成されています。したがって、F0、F1、F2、および F3 の 4 つの順方向パスと、B3、B2、B1、および B0 の逆方向パスがあります。

PP では、と呼ばれる、調整するための新しいハイパーパラメータが導入されています 块 (chunks)。同じパイプレベルを通じて連続して送信されるデータブロックの数を定義します。たとえば、図の下半分には、が表示されます chunks = 4。GPU0 は、チャンク 0、1、2、および 3 (F0,0、F0,1、F0,2、F0,3) で同じ順方向パスを実行し、他の GPU が作業を完了するまで待機してから、GPU0 が再び動作し始めます。ブロック 3、2、1、0 の逆方向パス (B0,3、B0,2、B0,1、B0,0)。

これは概念的には勾配累積ステップ (GAS) と同じであることに注意してください。PyTorch はそれを呼び出し 块、DeepSpeed はそれを呼び出します GAS。

なぜなら 块、PP ではマイクロバッチ (MBS) の概念が導入されているからです。DP はグローバルバッチサイズを小さなバッチサイズに分割するため、DP 次数が 4 の場合、グローバルバッチサイズ 1024 は 4 つの小さなバッチサイズに分割され、それぞれの小さなバッチサイズは 256 (1024/4) になります。数値 (または GAS) が 32 の場合 块 、マイクロバッチサイズは 8 (256/32) になります。各チューブステージは一度に 1 つのマイクロバッチを処理します。

DP + PP 設定のグローバルバッチサイズを計算する式は次のとおりです: mbs*chunks*dp_degree ( 8*32*4=1024)。

戻ってもう一度写真を見てみましょう。

最終的に得られたものを使用するの chunks=1 は単純な PP であり、非常に非効率的です。また 块 、数が非常に大きい場合、マイクロバッチサイズが小さくなり、おそらくあまり効率的ではありません。したがって、 块 GPU を最も効率的に使用できる数値を見つけるために実験する必要があります。

forward グラフは、最後のステージが backward パイプラインの終了を待つ必要があるため、並列化できない「デッド」タイムのバブルがあることを示しています。次に、参加しているすべての GPU が高い同時使用率を達成できるように最適な 块 数を見つけるという問題は、実際にはバブルの数を最小限に抑えることに変換されます。

このスケジューリングメカニズムはと呼ばれます 全前全后。その他のオプションとしては、タンデムおよびスタッガードタンデムがあります。

Megatron-LM と DeepSpeed は両方とも PP プロトコルの独自の実装を持っていますが、Megatron-DeepSpeed は DeepSpeed の他の機能と統合されているため、DeepSpeed 実装を使用します。

ここでのもう 1 つの重要な問題は、単語埋め込み行列のサイズです。一般に、単語埋め込み行列はトランスフォーマーブロックよりも少ないメモリを必要としますが、250,000 語彙を持つ BLOOM の場合、トランスフォーマーブロックの場合はわずか 4.9GB であるのに対し、埋め込み層は bf16 の重みに対して 7.2GB を必要とします。したがって、Megatron-Deepspeed に埋め込み層をトランスブロックとして扱わせる必要がありました。したがって、72 ステージのパイプラインがあり、そのうち 2 ステージ (最初と最後) は埋め込み専用です。これにより、GPU のメモリ消費のバランスをとることができます。これを行わなかった場合、最初と最後のステージで大量の GPU メモリが消費され、GPU メモリの使用量の 95% は非常にわずかになるため、トレーニングは非常に非効率になります。

DP+PP

DeepSpeed Pipeline Parallel チュートリアルには、以下に示すように、DP と PP を組み合わせる方法を示す図があります。

ここで理解しておくべき重要なことは、DP ランク 0 は GPU2 を認識できず、DP ランク 1 は GPU3 を認識できないことです。DP の場合、GPU 0 と 1 のみがあり、データはそれらに供給されます。GPU0 は PP を使用して、負荷の一部を GPU2 に「密かに」オフロードします。同様に、GPU1 も GPU3 からサポートを受けます。

各次元には少なくとも 2 つの GPU が必要なので、ここでは少なくとも 4 つの GPU が必要です。

DP+PP+TP

より効率的なトレーニングを行うには、下の図に示すように、3D 並列処理と呼ばれる PP、TP、および DP を組み合わせることができます。

この図はブログ投稿「3D Parallelism: Scaling to Trillion Parameter Models」) からのもので、これも優れた記事です。

次元ごとに少なくとも 2 つの GPU が必要なので、完全な 3D 並列処理には少なくとも 8 つの GPU が必要です。

ゼロDP+PP+TP

DeepSpeed の主な機能の 1 つは、ZeRO です。これは、「ZeRO データ並列処理」セクションで説明した、DP の超スケーラブルな拡張バージョンです。通常、これは独立した機能であり、PP や TP は必要ありません。ただし、PP、TPと組み合わせることもできます。

ZeRO-DP を PP (したがって TP) と組み合わせると、通常はオプティマイザー状態のシャーディングのみを行う ZeRO フェーズ 1 のみが有効になります。ZeRO ステージ 2 では勾配もシャーディングされ、ステージ 3 ではモデルの重みもシャーディングされます。

理論的にはパイプライン並列処理で ZeRO ステージ 2 を使用することは可能ですが、パフォーマンスに悪影響を与える可能性があります。各マイクロバッチでは、シャーディングの前に勾配を集約するための追加のリデュース散乱通信が必要ですが、これにより、重大な通信オーバーヘッドが追加される可能性があります。パイプラインの並列性に従って、小さなマイクロバッチを使用し、演算強度 (マイクロバッチサイズ) とパイプラインバブルの最小化 (マイクロバッチの数) の間のトレードオフに焦点を当てます。したがって、通信オーバーヘッドの増加により、パイプラインの並列処理が損なわれます。

また、PPのせいでレイヤーの数が通常よりも少ないため、あまりメモリを節約できません。PP は勾配サイズを縮小している 1/PPため、これに基づく勾配スライスは純粋な DP と比較して多くのメモリを節約しません。

ZeRO ステージ 3 を使用してこのサイズのモデルをトレーニングすることもできますが、並行して DeepSpeed 3D よりも多くの通信が必要になります。1 年前、私たちの環境を注意深く評価した結果、Megatron-DeepSpeed 3D 並列処理が最高のパフォーマンスを発揮することがわかりました。ZeRO Phase 3 のパフォーマンスは当時より大幅に向上しており、今再評価するとしたら、おそらく Phase 3 を選択するでしょう。

BF16オプティマイザー

FP16 を使用して巨大な LLM モデルをトレーニングすることは禁止されています。

私たちは104B モデルのトレーニングに何ヶ月も費やしてこれを自分たちで実証しましたが、 Tensorboardからわかるように、これは完全な失敗でした。絶え間なく発散する lm 損失と戦う過程で、私たちは多くのことを学びました。

また、530B モデルをトレーニングした後、Megatron-LM チームと DeepSpeed チームからも同じ提案を受けました。最近発売された OPT-175B もFP16で猛トレーニングを積んだと報告している。

そこで 1 月に、BF16 フォーマットをサポートする A100 でトレーニングすることを知っていました。Olatunji Ruwase は BLOOM のトレーニング用に開発しました BF16Optimizer。

このデータ形式に詳しくない場合は、そのビットレイアウトを見てください。BF16 形式の重要な点は、指数の数が FP32 と同じであるため、オーバーフローしないことですが、FP16 は頻繁にオーバーフローします。FP16 の最大値範囲は 64k で、より小さい数値のみを乗算できます。たとえば 250*250=62500、を行うことはできますが、を実行しようとすると 255*255=65025オーバーフローが発生し、これがトレーニングの問題の主な原因です。これは、重みを小さく保つ必要があることを意味します。ロススケーリングと呼ばれる手法はこの問題の軽減に役立ちますが、FP16 の範囲が狭いため、モデルが非常に大きくなる場合には依然として問題が発生する可能性があります。

BF16 にはこの問題はなく、簡単に実行できます 10_000*10_000=100_000_000。まったく問題ありません。

もちろん、BF16 と FP16 は同じサイズの 2 バイトであるため、フリーランチはなく、BF16 を使用する場合のトレードオフとして、精度が非常に低いことが挙げられます。ただし、トレーニングで使用した確率的勾配降下法とそのバリエーションは、スタッガリングに似ていることを覚えておいてください。このステップで完璧な方向が見つからなくても、問題はありません。次のステップで修正します。自分自身のステップ。

BF16 を使用するか FP16 を使用するかに関係なく、FP32 には常に重みのコピーが存在します。これはオプティマイザによって更新されるものです。したがって、16 ビット形式は計算にのみ使用され、オプティマイザは FP32 重みを完全な精度で更新し、次の反復でそれらを 16 ビット形式に変換します。

すべての PyTorch コンポーネントは、FP32 で蓄積を確実に実行できるように更新されているため、精度の損失は発生しません。

主要な問題は勾配の累積です。これは、各マイクロバッチによって処理される勾配が累積されるため、パイプライン並列処理の主な機能の 1 つです。トレーニングの精度を高めるために FP32 に勾配累積を実装することは重要であり、これがまさに BF16Optimizer 行われたことです。

他の改善点の中でも、次の lm 損失プロットに見られるように、BF16 混合精度トレーニングを使用することで、潜在的な悪夢が比較的スムーズなプロセスに変わったと考えられます。

CUDAフュージョンカーネル機能

GPU は主に 2 つのことを行います。ビデオメモリへのデータの書き込みとビデオメモリからのデータの読み取り、およびそのデータの計算を実行できます。GPU がデータの読み取りと書き込みでビジー状態であるとき、GPU のコンピューティングユニットはアイドル状態になります。GPU を効率的に利用したい場合は、アイドル時間を最小限に抑えたいと考えます。

カーネル関数は、特定の PyTorch 操作を実装する一連の命令です。たとえば、これを呼び出すと、 PyTorch スケジューラをtorch.add 通過し、入力テンソルやその他の変数の値に基づいて実行するコードを決定し、最終的に実行します。CUDA カーネルは CUDA を使用してこれらのコードを実装するため、NVIDIA GPU でのみ実行されます。

さて、 GPU で計算する場合、通常、PyTorch は 2 つの別々のカーネルを起動します。1 つは合計の c = torch.add (a, b); e = torch.max ([c,d]) 加算を実行し、もう 1 つは 2 つの最大値を取得します。この場合、GPU はビデオメモリから合計をフェッチし、加算を実行して、結果をビデオメモリに書き込みます。次に、合計を取得して最大演算を実行し、結果をビデオメモリに再度書き込みます。abcdabcd

これら 2 つの操作を融合する場合、つまり、それらを「融合されたカーネル関数」に入れてから、中間結果をビデオメモリに書き込む代わりに、そのカーネルを起動する場合、中間結果を GPU レジスタに保持し、Get を実行するだけで済み c ます d 。最終的な計算。これにより、オーバーヘッドが大幅に節約され、GPU のアイドリングが防止されるため、操作全体がはるかに効率的になります。

フュージョンカーネル関数はまさにそれを行います。これらは主に、複数の離散計算とビデオメモリとの間のデータ移動を、データ移動が非常に少ない融合計算に置き換えます。さらに、一部のフュージョンカーネルは、演算を数学的に変換して、特定の計算の組み合わせをより高速に実行できるようにします。

BLOOM を迅速かつ効率的にトレーニングするには、Megatron-LM が提供するいくつかのカスタム CUDA 融合カーネル関数を使用する必要があります。特に、LayerNorm フュージョンカーネルと、フュージョンスケーリング、マスキング、ソフトマックス操作のさまざまな組み合わせ用のカーネルがあります。Bias Add は、PyTorch の JIT 機能を通じて GeLU とも統合されます。これらの操作はすべてメモリに依存するため、ビデオメモリの各読み取り後の計算量を最大化するためにそれらを融合することが重要です。したがって、たとえば、メモリ内にボトルネックがある GeLU 操作の実行中に Bias Add を実行しても、実行時間は増加しません。これらのカーネル関数は、Megatron-LM リポジトリコードベースにあります。

データセット

Megatron-LM のもう 1 つの重要な機能は、効率的なデータローダーです。最初のトレーニングを開始する前に、各データセットの各サンプルは固定シーケンス長 (BLOOM は 2048) のサンプルに分割され、各サンプルに番号を付けるインデックスが作成されます。トレーニングハイパーパラメーターに基づいて、各データセットが参加する必要があるエポックの数を決定し、これに基づいてサンプルインデックスの順序付きリストを作成し、それをシャッフルします。たとえば、データセットに 2 エポックでトレーニングする必要がある 10 個のサンプルがある場合、システムはまずサンプルインデックスを順番に並べ替え [0, ..., 9, 0, ..., 9] 、次に順序をシャッフルしてデータセットの最終的なグローバル順序を作成します。これは、トレーニングが単にデータセット全体を反復処理して繰り返すわけではないことを意味することに注意してください。同じサンプルが 2 回表示されてから、別のサンプルが表示される可能性がありますが、トレーニングの最後には、モデルは各サンプルを 2 回だけ表示します (セカンドレート)。これにより、トレーニング全体を通してスムーズなトレーニング曲線を確保できます。元のデータセット内の各サンプルのオフセットを含むこれらのインデックスは、トレーニングが開始されるたびに再計算されることを避けるためにファイルに保存されます。最後に、これらのデータセットのいくつかを異なる重みでブレンドして、トレーニングに使用される最終データを作成できます。

埋め込みレイヤーノルム

104B モデルの発散を防ぐ取り組みの中で、最初の単語埋め込み層の後に LayerNorm を追加すると、トレーニングがより安定することがわかりました。

この洞察は、StableEmbedding 一律の xavier 関数で初期化された LayerNorm を使用した通常の埋め込みである操作を持つbitsandbytes を使用した実験から得られます。

場所コード

「短くトレーニングし、長くテスト: 線形バイアスによる注意が入力長の外挿を可能にする」という論文に基づいて、通常の位置埋め込みを AliBi に置き換えます。これにより、モデルのトレーニングに使用される入力シーケンスよりも長い入力シーケンスの外挿が可能になります。したがって、長さ 2048 のシーケンスでトレーニングしても、モデルは推論中により長いシーケンスを処理できます。

トレーニングの難しさ

アーキテクチャ、ハードウェア、ソフトウェアが整備されたことで、2022 年 3 月初旬にトレーニングを開始することができました。しかし、それ以来、すべてが順風満帆だったわけではありません。このセクションでは、私たちが遭遇した主な障害のいくつかについて説明します。

トレーニングを開始する前に、解決すべき質問がたくさんあります。特に、小規模ではなく 48 ノードでトレーニングを開始した後にのみ現れる問題がいくつか見つかりました。たとえば、 CUDA_LAUNCH_BLOCKING=1 フレームワークがハングしないようにするには、オプティマイザグループをより小さなグループに分割する必要があります。そうしないと、フレームワークが再びハングしてしまいます。これらについて詳しくは、トレーニング前の記録をご覧ください。

トレーニング中に発生する主な問題はハードウェア障害です。これは約 400 個の GPU を備えた新しいクラスターであるため、平均して 1 週間に 1 ～ 2 個の GPU 障害が発生します。3 時間ごとにチェックポイントを保存します (100 回の反復)。その結果、ハードウェアのクラッシュにより、週に平均 1.5 時間のトレーニングが失われます。Jean Zay システム管理者は、障害のある GPU を交換し、ノードを復元します。それまでの間、予備のノードをご用意しております。

他にもさまざまな問題が発生し、5 ～ 10 時間のダウンタイムが複数回発生しました。その一部は PyTorch のデッドロックバグに関連しており、その他はディスク領域の不足が原因でした。詳細に興味がある場合は、トレーニング記録を参照してください。

このダウンタイムはすべて、このモデルのトレーニングの実現可能性分析で計画されており、適切なモデルサイズと、それに応じてモデルが消費するデータ量を選択しました。そのため、こうしたダウンタイムの問題はありましたが、なんとか推定時間内にトレーニングを完了することができました。前述したように、完了までに約 100 万コンピューティング時間かかります。

もう 1 つの問題は、SLURM が人々のグループによって使用されるように設計されていないことです。SLURM ジョブは 1 人のユーザーによって所有され、そのユーザーがいない場合、グループの他のメンバーは実行中のジョブに対して何も行うことができません。プロセスを開始したユーザーがいなくても、グループ内の他のユーザーが現在のプロセスを終了できる終了スキームがあります。これは問題の 90% に対してうまく機能します。SLURM 設計者がこれを読んだ場合は、SLURM ジョブをグループで所有できるように Unix グループの概念を追加してください。

トレーニングは 24 時間年中無休で行われるため、待機中の人が必要ですが、ヨーロッパとカナダの西海岸に人がいるため、誰かがポケベルを持ち歩く必要はなく、お互いをサポートするのが得意です。もちろん、週末のトレーニングには注意が必要です。ハードウェアクラッシュからの自動回復など、ほとんどのことを自動化していますが、それでも人間の介入が必要な場合があります。

重要なリンク

論文と記事

この記事ですべてを詳しく説明するのは不可能なので、ここで紹介したテクニックが好奇心をそそり、さらに詳しく知りたい場合は、次の論文を読んでください。

メガトロン-LM:

ディープスピード:

Megatron-LM と Deepspeedeed の組み合わせ:

DeepSpeed と Megatron を使用した Megatron-Turing NLG 530B、大規模生成言語モデルのトレーニング。

アリバイ：

トレーニングは短く、テストは長時間: 線形バイアスによる注意により、入力長の外挿が可能
100 万 GPU 時間がある場合、どの言語モデルをトレーニングするか? - 最終的に ALiBi を選択するきっかけとなった実験がここにあります。

ビットNバイト:

ブロック単位の量子化による 8 ビットオプティマイザー (この論文では埋め込み LaynerNorm を使用しましたが、論文の他の部分とそのテクニックも非常に優れています。8 ビットオプティマイザーを使用しなかった唯一の理由は、すでに使用していたということです。 DeepSpeed-ZeRO オプティマイザーメモリを節約します)。

Megatron-DeepSpeed についての一般的な理解: 1,000 億のパラメーターモデル BLOOM を支えるテクノロジー

序文