BLOOMトレーニングの技術背後

近年、ますます大規模な言語モデルの訓練が一般的になってきました。これらのモデルがさらなる研究のために公開されていない問題は頻繁に議論されますが、そのようなモデルを訓練するための技術やエンジニアリングについての隠された知識は滅多に注目されません。本記事では、1760億パラメータの言語モデルBLOOMを例に、そのようなモデルの訓練の裏側にあるハードウェアとソフトウェアの技術とエンジニアリングについて、いくつかの光を当てることを目指しています。

しかし、まず、この素晴らしい1760億パラメータモデルの訓練を可能にするために貢献してくれた企業や主要な人物やグループに感謝したいと思います。

その後、ハードウェアのセットアップと主要な技術的な構成要素について説明します。

以下はプロジェクトの要約です：

人々

このプロジェクトは、Hugging Faceの共同創設者でありCSOのThomas Wolf氏が考案しました。彼は巨大な企業と競争し、単なる夢だったものを実現し、最終的な結果をすべての人にアクセス可能にすることで、最も多くの人々にとっては夢であったものを実現しました。

この記事では、モデルの訓練のエンジニアリング側に特化しています。BLOOMの背後にある技術の最も重要な部分は、私たちにコーディングと訓練の助けを提供してくれた専門家の人々と企業です。

感謝すべき6つの主要なグループがあります：

HuggingFaceのBigScienceチームは、数人の専任の従業員を捧げ、訓練を始めから終わりまで行うための方法を見つけるために、Jean Zayの計算機を超えるすべてのインフラストラクチャを提供しました。
MicrosoftのDeepSpeedチームは、DeepSpeedを開発し、後にMegatron-LMと統合しました。彼らの開発者たちはプロジェクトのニーズに多くの時間を費やし、訓練前後に素晴らしい実践的なアドバイスを提供しました。
NVIDIAのMegatron-LMチームは、Megatron-LMを開発し、私たちの多くの質問に親切に答えてくれ、一流の実践的なアドバイスを提供しました。
ジャン・ゼイのスーパーコンピュータを管理しているIDRIS / GENCIチームは、計算リソースをプロジェクトに寄付し、優れたシステム管理のサポートを提供しました。
PyTorchチームは、このプロジェクトのために基礎となる非常に強力なフレームワークを作成し、訓練の準備中に私たちをサポートし、複数のバグを修正し、PyTorchコンポーネントの使いやすさを向上させました。
BigScience Engineeringワーキンググループのボランティア

プロジェクトのエンジニアリング側に貢献してくれたすべての素晴らしい人々を全て挙げることは非常に困難なので、Hugging Face以外のいくつかの主要な人物を挙げます。彼らはこのプロジェクトのエンジニアリングの基盤となりました。

Olatunji Ruwase、Deepak Narayanan、Jeff Rasley、Jared Casper、Samyam Rajbhandari、Rémi Lacroix

また、このプロジェクトに貢献するために従業員に貢献を許可したすべての企業に感謝しています。

概要

BLOOMのアーキテクチャは、GPT3と非常に似ており、後ほどこの記事で詳しく説明されます。

モデルは、フランス政府が資金提供し、フランス国立科学研究センター（CNRS）の国立計算センターであるIDRISにインストールされているGENCIが管理するJean Zayで訓練されました。計算リソースはGENCIによって寛大にプロジェクトに寄付されました（助成金2021-A0101012475）。

訓練中には以下のハードウェアが使用されました：

GPU：384台のNVIDIA A100 80GB GPU（48ノード）+ 32台の予備GPU
ノードごとに8つのGPUを使用し、NVLink 4のインターグプ接続、4つのOmniPathリンクを使用
CPU：AMD EPYC 7543 32コアプロセッサ
CPUメモリ：ノードごとに512GB
GPUメモリ：ノードごとに640GB
ノード間接続：Omni-Path Architecture（OPA）（ノンブロッキングファットツリー）
NCCL通信ネットワーク：完全に専用のサブネット
ディスクIOネットワーク：他のノードやユーザーと共有のGPFS

チェックポイント：

メインのチェックポイント
fp32最適化状態とbf16 + fp32ウェイトを持つ各チェックポイントのサイズは2.3TBです- bf16ウェイトだけでも329GBです。

データセット：

1.5TBの中で重複排除された大量のクリーンアップされたテキストの46言語、それを350Bのユニークなトークンに変換しました
モデルの語彙サイズは250,680トークンです
詳細については、「The BigScience Corpus A 1.6TB Composite Multilingual Dataset」をご覧ください

176B BLOOMモデルのトレーニングは2022年3月から7月まで行われ、約3.5ヶ月かかりました（約100万コンピュート時間）。

Megatron-DeepSpeed

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

DeepSpeedは、分散トレーニングを簡単で効率的、効果的に行うための深層学習最適化ライブラリです。
Megatron-LMは、NVIDIAの応用深層学習研究チームによって開発された大規模で強力なトランスフォーマーモデルフレームワークです。

DeepSpeedチームは、DeepSpeedライブラリのZeROシャーディングとパイプライン並列処理、およびMegatron-LMのテンソル並列処理を組み合わせた3D並列処理ベースの実装を開発しました。各コンポーネントの詳細は、以下の表で確認できます。

BigScienceのMegatron-DeepSpeedは、元のMegatron-DeepSpeedリポジトリのフォークであり、複数の追加機能を追加しました。

以下は、どのフレームワークがBLOOMのトレーニングに提供されたコンポーネントの表です：

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

Megatron-DeepSpeedは、3D並列処理を実装して、巨大なモデルを非常に効率的にトレーニングできるようにしています。3Dのコンポーネントについて簡単に説明しましょう。

DataParallel（DP） – 同じセットアップが複数回複製され、それぞれがデータのスライスを受け取ります。処理は並列で行われ、すべてのセットアップは各トレーニングステップの最後に同期されます。
TensorParallel（TP） – 各テンソルは複数のチャンクに分割されるため、テンソル全体が単一のGPU上に存在するのではなく、テンソルのシャードがそれぞれの指定されたGPU上に配置されます。処理中には、各シャードが別々のGPU上で並列に処理され、結果はステップの最後で同期されます。これは、水平並列処理と呼ぶことができるもので、分割は水平レベルで行われます。
PipelineParallel（PP） – モデルが垂直（層レベル）に複数のGPUに分割されるため、モデルの1つまたは複数の層が単一のGPUに配置されます。各GPUはパイプラインの異なるステージを並列に処理し、バッチの小さなチャンクで作業します。
Zero Redundancy Optimizer（ZeRO） – TPと似たようなテンソルのシャーディングも行いますが、全体のテンソルは順方向または逆方向の計算のために時間内に再構築されるため、モデルを変更する必要はありません。また、GPUメモリの制約に対応するためのさまざまなオフロード技術もサポートしています。

データ並列処理

数台のGPUを持つほとんどのユーザーは、DistributedDataParallel（DDP）PyTorchドキュメントに詳しいはずです。この方法では、モデルは各GPUに完全に複製され、各反復後にすべてのモデルが互いの状態を同期します。このアプローチは、問題により多くのリソースを投入することでトレーニング速度を向上させるが、モデルが単一のGPUに収まる場合にしか機能しません。

ZeROデータ並列処理

ZeROパワードデータ並列処理（ZeRO-DP）は、このブログ記事の図から説明されています

それを理解するのは難しいかもしれませんが、実際には、コンセプトは非常にシンプルです。これは通常のDDPと同じですが、フルモデルのパラメータ、勾配、およびオプティマイザの状態を完全に複製する代わりに、各GPUはそれぞれのスライスのみを保存します。そして、実行時には、与えられたレイヤーに対して必要な完全なレイヤーパラメータが必要な場合には、すべてのGPUが互いに欠落している部分を提供するために同期します – これだけです。

このコンポーネントはDeepSpeedによって実装されています。

テンソル並列処理

テンソル並列処理（TP）では、各GPUはテンソルのスライスのみを処理し、全体のテンソルは全体を必要とする操作のためにのみ集約されます。

このセクションでは、Megatron-LMの論文からの概念と図を使用します：GPUクラスタ上での効率的な大規模言語モデルのトレーニング。

トランスフォーマーの主要な構成要素は、完全に接続されたnn.Linearと、非線形活性化GeLUです。

Megatronの論文の表記法に従って、ドット積の一部をY = GeLU(XA)と書くことができます。ここで、XとYは入力と出力ベクトルであり、Aは重み行列です。

行列形式で計算を見ると、行列の乗算を複数のGPUで分割する方法がわかります：

重み行列AをN個のGPUに列方向に分割し、並列で行列乗算XA_1からXA_nを行うと、N個の出力ベクトルY_1, Y_2, ..., Y_nが得られます。これらは独立してGeLUに入力することができます：。Y行列が列方向に分割されていることに注意してください。2番目のGEMMも行方向に分割することができますので、余分な通信なしでGeLUの出力を直接取得することができます。

この原理を使用して、任意の深さのMLPを更新することができます。ただし、各行列乗算の後にGPUを同期する必要があります。Megatron-LMの著者はそのための役立つイラストを提供しています：

ここで、fはフォワードパスでの恒等演算子であり、バックワードパスではすべてのリデュースを行います。一方、gはフォワードパスでのリデュース演算子であり、バックワードパスでは恒等演算子です。

マルチヘッドアテンションレイヤーの並列化はさらに簡単です。複数の独立したヘッドを持つため、既に固有の並列性を持っています！

特別な考慮事項：TPでは、各層でのフォワードパスとバックワードパスの両方で2つのリデュースが必要なため、デバイス間の非常に高速なインターコネクトが必要です。そのため、非常に高速なネットワークがない限り、1つのノードを超えてTPを行うことはお勧めできません。私たちの場合、ノード間の通信はPCIeよりもはるかに遅かったです。実際には、ノードに4つのGPUがある場合、最大のTP度数は4です。TP度数が8が必要な場合、少なくとも8つのGPUを持つノードを使用する必要があります。

このコンポーネントはMegatron-LMによって実装されています。Megatron-LMは最近、シーケンスの次元に沿って分割できない操作（例：LayerNorm）を含むシーケンス並列化を追加しました。詳細については、「大規模トランスフォーマーモデルにおける活性化再計算の削減」の論文を参照してください。シーケンス並列化は、BLOOMのトレーニング前に開発されたため、BLOOMのトレーニングでは使用されません。

パイプライン並列化

ナイーブなパイプライン並列化（ナイーブ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に戻す必要があります（または、ラベルを最後のレイヤーに送信することもできます）。これで損失を計算し、最適化プログラムが作業を行うことができます。

問題点：

メインの欠点であり、なぜこれが「素朴な」PPと呼ばれているかは、1つを除いてすべてのGPUが任意の時点でアイドル状態になっていることです。したがって、4つのGPUを使用する場合、1つのGPUのメモリ量を4倍にするのとほぼ同じであり、その他のハードウェアは無視されます。さらに、デバイス間でデータをコピーするオーバーヘッドがあります。したがって、素朴なPPを使用すると、4つの6GBカードは1つの24GBカードとほぼ同じサイズを収容できますが、データコピーのオーバーヘッドがないため、後者の方がトレーニングが速く完了します。ただし、たとえば40GBのカードを持ち、45GBのモデルを収める必要がある場合、4つの40GBカード（ただし、勾配と最適化プログラムの状態のためぎりぎり）で収めることができます。
共有埋め込みは、GPU間でコピーする必要がある場合があります。

パイプライン並列処理（PP）は、上記で説明した素朴なPPとほぼ同じですが、GPUのアイドリング問題を解決するために、バッチをマイクロバッチに分割し、パイプラインを人工的に作成することにより、さまざまなGPUが同時に計算プロセスに参加できるようにします。

以下のGPipe論文の図は、上部に素朴なPP、下部にPPを示しています：

下の図からは、PPの方がアイドル状態である部分が少ないことがわかります。アイドル状態の部分は「バブル」と呼ばれます。

図の両方の部分は、パイプラインに参加する4つのGPUの並列処理を示しています。つまり、4つのパイプステージF0、F1、F2、F3の順方向パスと、B3、B2、B1、B0の逆順の逆方向パスがあります。

PPは、chunksと呼ばれる新しいハイパーパラメータを導入します。これは、データのチャンクが同じパイプステージを通じて連続して送信される回数を定義します。たとえば、下の図では、chunks=4となっています。GPU0は、チャンク0、1、2、3（F0,0、F0,1、F0,2、F0,3）で同じ順方向パスを実行し、他のGPUが作業を開始し始めてからのみ、逆方向パスのチャンク3、2、1、0（B0,3、B0,2、B0,1、B0,0）を実行します。

概念的には、これは勾配蓄積ステップ（GAS）と同じ概念です。PyTorchではchunksを使用しますが、DeepSpeedでは同じハイパーパラメータをGASと呼びます。

チャンクのため、PPはマイクロバッチ（MBS）の概念を導入します。DPはグローバルデータバッチサイズをミニバッチに分割するため、DPの度数が4である場合、グローバルバッチサイズ1024は4つのミニバッチ（各256）に分割されます（1024/4）。そして、chunks（またはGAS）の数が32である場合、マイクロバッチサイズは8になります（256/32）。各パイプラインステージは、一度に1つのマイクロバッチで作業します。

DP + PPのセットアップのグローバルバッチサイズを計算するには、mbs*chunks*dp_degree（8*32*4=1024）を行います。

図に戻りましょう。

chunks=1の場合、非効率な素朴なPPになります。非常に大きなchunks値の場合、非常に効率的でない小さなマイクロバッチサイズになります。したがって、最も効率的なGPUの利用を実現する値を見つけるためには、実験を行う必要があります。

図では、最後のforwardステージがbackwardの完了を待たなければならないため、並列化できない「デッド」時間のバブルがあることが示されています。しかし、chunksの最適な値を見つけることで、すべての参加するGPUを高い並行利用率で利用できるようにし、バブルのサイズを最小限に抑えることができます。

このスケジューリングメカニズムは、all forward all backward として知られています。他の代替手段としては、one forward one backwardやinterleaved one forward one backwardなどがあります。

Megatron-LMとDeepSpeedの両方には、PPプロトコルの独自の実装がありますが、Megatron-DeepSpeedではDeepSpeedの実装が使用されています。DeepSpeedとの他の側面と統合されているためです。

ここでのもう1つの重要な問題は、単語埋め込み行列のサイズです。通常、単語埋め込み行列はトランスフォーマーブロックよりも少ないメモリを消費しますが、巨大な25万語彙の場合、埋め込み層には7.2GBのbf16ウェイトが必要であり、トランスフォーマーブロックはわずか4.9GBです。そのため、Megatron-Deepspeedに埋め込み層をトランスフォーマーブロックとして扱うよう指示する必要がありました。そのため、埋め込み（最初と最後）に専用の2つのレイヤーを含む72層のパイプラインがありました。これにより、GPUメモリの消費量をバランスさせることができました。これを行わなかった場合、最初と最後のステージがほとんどのGPUメモリを消費し、95％のGPUがメモリをはるかに少なく使用しているため、トレーニングは効率的ではありませんでした。

DP+PP

DeepSpeedパイプラインチュートリアルの次の図は、DPとPPを組み合わせる方法を示しています。

ここで重要なのは、DPランク0はGPU2を見ることができず、DPランク1はGPU3を見ることができないということです。DPにとって、データが2つのGPUだけあるかのようにデータを供給しています。GPU0はPPを使用してGPU2に一部の負荷を「秘密裏に」オフロードし、GPU1もGPU3を助けに立てることで同様のことをします。

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

DP+PP+TP

より効率的なトレーニングを行うために、PPとTPとDPを組み合わせた3D並列処理と呼ばれる手法があります。次の図に示されています。

この図は、トリリオンパラメータモデルにスケーリングするための3D並列処理に関するブログ記事「3D parallelism: Scaling to trillion-parameter models」から取得されたもので、おすすめの読み物です。

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

ZeRO DP+PP+TP

DeepSpeedの主な特徴の1つは、DPの超スケーラブルな拡張であるZeROです。これはすでに「ZeRO Data Parallelism」で説明されています。通常、これはPPやTPを必要としないスタンドアロンの機能ですが、PPやTPと組み合わせることもできます。

ZeRO-DPがPPと組み合わされる場合（必要に応じてTPも）、通常はZeROステージ1のみが有効になります。これはオプティマイザの状態のみをシャードします。ZeROステージ2では、勾配もシャードし、ステージ3ではモデルの重みもシャードします。

PPにPipeline Parallelismを組み合わせてZeROステージ2を使用することは理論的には可能ですが、パフォーマンスに悪影響を与えます。各マイクロバッチごとに追加のリデューススキャッタコレクティブが必要になり、勾配をシャードする前に勾配を集約するため、潜在的に重要な通信オーバーヘッドが発生します。Pipeline Parallelismの性質上、小さなマイクロバッチが使用され、マイクロバッチの数を最小化することに重点が置かれます。したがって、そのような通信コストは悪影響を及ぼすでしょう。

さらに、PPによって通常よりもレイヤーが少なくなっているため、メモリの節約はそれほど大きくありません。PPは既に勾配サイズを1/PP削減しているため、それに加えて勾配シャーディングの節約は純粋なDPよりも重要ではありません。

ZeROステージ3もこのスケールでモデルをトレーニングするために使用することができますが、DeepSpeed 3Dパラレル実装よりも通信が必要です。1年前に環境で慎重に評価した結果、Megatron-DeepSpeed 3D並列処理が最も優れたパフォーマンスを発揮しました。その後、ZeROステージ3のパフォーマンスは大幅に改善されており、今日評価する場合、おそらくステージ3を選択していたでしょう。

BF16Optimizer

大規模なLLMモデルをFP16でトレーニングすることは避けるべきです。

私たちは、数ヶ月にわたって104Bモデルのトレーニングを行い、テンソルボードからわかるように完全な失敗であったことを証明しました。常に発散するlm-lossと戦いながら、多くのことを学びました。

Megatron-LMとDeepSpeedのチームも530Bモデルをトレーニングした後、同じアドバイスを受けました。OPT-175Bの最新リリースでも、FP16でのトレーニングが非常に難しかったと報告されています。

そのため、私たちはA100でトレーニングすることを知っていたため、BF16形式をサポートする「BF16Optimizer」を開発し、それを使用してBLOOMをトレーニングしました。

このデータフォーマットに慣れていない場合は、ビットのレイアウトをご覧ください。BF16形式のキーは、FP32と同じ指数を持っているため、FP16が頻繁に発生するオーバーフローの問題に悩まされません！FP16では、最大数値範囲が64kであり、小さな数値のみを乗算できます。例えば、250*250=62500はできますが、255*255=65025を試すとオーバーフローが発生し、トレーニング中の主な問題の原因となります。つまり、重みは小さくなければなりません。損失スケーリングという技術は、この問題に対処するのに役立つかもしれませんが、FP16の限られた範囲は、モデルが非常に大きくなると依然として問題です。

BF16にはこのような問題はありません。簡単に10_000*10_000=100_000_000などができます。

もちろん、BF16とFP16は2バイトのサイズが同じであるため、完全な精度を犠牲にすることになります。ただし、確率的勾配降下法とその変種を使用したトレーニングは、つまづくような歩行のようなものです。したがって、最初に完璧な方向を得なくても問題ありません。次のステップで修正します。

BF16またはFP16を使用するかどうかに関係なく、常にFP32の重みのコピーが存在します。これはオプティマイザによって更新されます。したがって、16ビットフォーマットは計算にのみ使用され、オプティマイザは完全な精度でFP32の重みを更新し、次のイテレーションのために16ビットフォーマットに変換します。

すべてのPyTorchコンポーネントは、蓄積をFP32で実行するように更新されており、そこで損失は発生しません。

重要な問題の1つは勾配の蓄積であり、各マイクロバッチ処理の勾配が蓄積されるパイプライン並列処理の主な特徴の1つです。トレーニングの精度を保つために、勾配蓄積をFP32で実装することが重要であり、これが「BF16Optimizer」が行うことです。

その他の改善策に加えて、BF16ミックスプレシジョントレーニングの使用により、潜在的な悪夢が比較的スムーズなプロセスに変わったと考えています。以下のlm損失グラフからもわかります。

Fused CUDAカーネル

GPUは2つのことを行います。メモリへのデータのコピーとそのデータに対する計算です。GPUがデータのコピーを行っている間、計算ユニットはアイドル状態になります。GPUを効率的に利用するためには、アイドル時間を最小限に抑える必要があります。

カーネルは、特定のPyTorch操作を実装する一連の命令です。たとえば、torch.addを呼び出すと、PyTorchディスパッチャが入力テンソルなどを確認し、実行するコードを決定し、実行します。CUDAカーネルは、CUDA APIライブラリを使用する特定の実装であり、NVIDIAのGPUでのみ実行できます。

例えば、GPUにc = torch.add(a, b); e = torch.max([c,d])の計算を指示する場合、単純なアプローチであり、PyTorchが指示されていない場合に行うことは、2つの別々のカーネルを起動することです。1つはaとbの加算を実行し、もう1つはcとdの間の最大値を求めます。この場合、GPUは自身のメモリからaとbを取得し、加算を実行し、その結果をメモリにコピーします。次にcとdを取得し、max操作を実行し、再び結果をメモリにコピーします。

もしもこれらの2つの操作を融合させ、つまりそれらを1つの「融合カーネル」としてまとめ、その1つのカーネルを起動するだけで、中間結果 c をメモリにコピーせず、GPUのレジスタに残し、最後の計算を完了するために d を取得する必要があります。これにより、多くのオーバーヘッドが削減され、GPUのアイドル状態が防止され、全体の操作がはるかに効率的になります。

融合カーネルはそのようなものです。主に、複数の個別の計算とメモリへのデータ移動を融合された計算に置き換え、非常に少数のメモリ移動を持つ融合計算にします。さらに、一部の融合カーネルは数式を書き換えて、特定の計算グループをより高速に実行できるようにします。

BLOOMを高速かつ効率的にトレーニングするために、Megatron-LMが提供するいくつかのカスタム融合CUDAカーネルを使用する必要がありました。特に、LayerNormを実行するための最適化されたカーネル、およびスケーリング、マスキング、およびソフトマックス操作のさまざまな組み合わせを融合させるためのカーネルがあります。バイアス項目もPyTorchのJIT機能を使用してGeLU操作と融合されます。これらの操作はすべてメモリに依存しているため、メモリから値が取得されると計算が行われる量を最大化するためにそれらを融合することが重要です。したがって、例えば、メモリに依存したGeLU操作を行っている間にバイアス項目を追加すると、追加の時間がかかりません。これらのカーネルはすべてMegatron-LMリポジトリで利用できます。

データセット

Megatron-LMのもう1つの重要な機能は、効率的なデータローダーです。初期トレーニングの開始時に、各データセットはリクエストされたシーケンス長（BLOOMの場合は2048）のサンプルに分割され、それぞれのサンプルに番号を付けるためのインデックスが作成されます。トレーニングパラメータに基づいて、データセットのエポック数が計算され、その数だけの順序が作成され、シャッフルされます。たとえば、データセットに10のサンプルがあり、2回通過する必要がある場合、システムは最初にサンプルのインデックスを[0、...、9、0、...、9]の順に配置し、その順序をシャッフルしてデータセットの最終的なグローバル順序を作成します。これにより、トレーニングは単純にデータセット全体を通過してから繰り返すだけではなく、同じサンプルを2回見ることがありますが、トレーニングの最後までにモデルは各サンプルを2回見ることになります。これにより、トレーニングプロセス全体でスムーズなトレーニングカーブが確保されます。これらのインデックス、および各サンプルの基本データセットへのオフセットも、トレーニングプロセスを開始するたびに再計算することを避けるために、ファイルに保存されます。これらのデータセットのいくつかは、トレーニングプロセスが見る最終的なデータに異なる重みを加えることでブレンドすることができます。

埋め込みレイヤーノルム

104Bの発散を止めようと奮闘している間に、最初の単語の埋め込みの直後に追加のレイヤーノルムを追加することで、トレーニングがはるかに安定することがわかりました。

この洞察は、bitsandbytesでの実験から得られました。bitsandbytesには通常の埋め込みにレイヤーノルムが含まれており、一様なxavier初期化が使用されています。

位置エンコーディング

通常の位置埋め込みを、AliBiに基づいたものに置き換えました。これは、Train Short、Test Long: Attention with Linear Biases Enables Input Length Extrapolationという論文に基づいています。これにより、モデルがトレーニングされたシーケンスよりも長い入力シーケンスに対して外挿することができます。したがって、2048の長さのシーケンスでトレーニングしていても、モデルは推論中にはるかに長いシーケンスにも対応できます。

トレーニングの困難

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

トレーニングを開始する前に解決する必要がある問題がたくさんありました。特に、48ノードでのトレーニングを開始した場合にのみ現れる問題がいくつかありました。たとえば、フレームワークがフリーズするのを防ぐためにCUDA_LAUNCH_BLOCKING=1が必要であり、オプティマイザのグループを小さなグループに分割する必要がありました。さもないと、フレームワークが再度フリーズしてしまいます。これについての詳細は、トレーニングプリクエルクロニクルで詳しく読むことができます。

トレーニング中に遭遇した主な問題のタイプはハードウェアの故障でした。これは約400のGPUを備えた新しいクラスタであったため、平均して週に1〜2回のGPUの故障が発生していました。私たちは3時間ごと（100イテレーションごと）にチェックポイントを保存していたため、平均して1.5時間のトレーニングがハードウェアのクラッシュで失われました。Jean Zayのシステム管理者は故障したGPUを交換し、ノードを再起動しました。その間、代わりにバックアップノードを使用しました。

私たちは、PyTorchのデッドロックバグやディスク容量の不足など、さまざまな問題に遭遇し、数回にわたって5〜10時間のダウンタイムが発生しました。具体的な詳細に興味がある場合は、トレーニングの記録をご覧ください。

このモデルのトレーニングの実現可能性を決定する際に、これらのダウンタイムをすべて予測して計画しました。私たちはモデルのサイズをその実現可能性とモデルが消費するデータ量に合わせて選びました。すべてのダウンタイムを経て、予想された時間内にトレーニングを完了することができました。先に述べたように、約100万計算時間がかかりました。

もう1つの問題は、SLURMがチームで使用するために設計されていなかったことです。SLURMのジョブは単一のユーザーが所有し、そのユーザーがいない場合は、グループの他のメンバーは実行中のジョブに何もできません。私たちは、起動したプロセスが存在しなくても、グループの他のユーザーが現在のプロセスを停止できるキルスイッチの回避策を開発しました。これは90％の問題でうまく機能しました。SLURMの設計者がこれを読んでいる場合は、SLURMジョブがグループによって所有される概念を追加してください。

トレーニングが24時間365日行われていたため、常に担当者が待機している必要がありました。ただし、ヨーロッパとカナダの西海岸に人がいたため、誰かがポケベルを持つ必要はありませんでした。私たちは単に上手く重なり合わせました。もちろん、週末には誰かがトレーニングを監視する必要がありました。私たちは、ハードウェアのクラッシュからの回復を含むほとんどのことを自動化しましたが、時には人間の介入も必要でした。

結論

トレーニングの最も困難で集中的な部分は、トレーニングの開始までの2か月間でした。リソースの割り当てが時間的に制限されており、A100にアクセスできるのは非常に最後の瞬間まででしたので、なるべく早くトレーニングを開始するというプレッシャーがありました。そのため、BF16Optimizerは最後の瞬間に書かれ、デバッグしてさまざまなバグを修正する必要がありました。また、前のセクションで説明したように、48ノードでトレーニングを開始した後に現れる新しい問題も発見しましたが、小規模なスケールでは現れません。

しかし、これらの問題を解決すると、トレーニングそのものは驚くほどスムーズで、大きな問題もありませんでした。ほとんどの時間は1人の担当者がトレーニングを監視し、数回にわたって複数の人々がトラブルシューティングに関与しました。Jean Zayの管理部門からは、トレーニング中に発生したほとんどのニーズに迅速に対応していただきました。

全体的には、非常に集中的でありながら非常に報酬のある経験でした。

大規模言語モデルのトレーニングはまだ課題がありますが、この技術をオープンに構築し共有することで、他の人々が私たちの経験に基づいてさらに進歩することを願っています。

リソース

重要なリンク

メインのトレーニングドキュメント
TensorBoard
トレーニング用のSLURMスクリプト
トレーニングの記録

論文と記事

この記事ではすべての詳細を詳しく説明することはできませんでしたので、ここで紹介された技術に興味がある場合、以下の論文を読んでください:

Megatron-LM:

GPUクラスターでの効率的な大規模言語モデルのトレーニング
大規模トランスフォーマーモデルにおける活性再計算の削減

DeepSpeed:

トリリオンパラメータモデルのトレーニングに向けたメモリ最適化
ビリオンスケールモデルトレーニングの民主化
極限スケールの深層学習におけるGPUメモリの壁を打破する
すべての人のための極限スケールモデルトレーニング

Megatron-LMとDeepSpeedの共同:

Megatron-Turing NLG 530BのトレーニングにDeepSpeedとMegatronを使用する

ALiBi:

Attention with Linear Biasesによるトレーニングの短縮、テストの延長
100万GPU時間がある場合にどの言語モデルをトレーニングするか？- ALiBiを選択するための実験結果が記載されています。

BitsNBytes:

8ビット最適化によるブロックごとの量子化（埋め込みLayerNormの文脈でのみですが、その他の部分と技術は素晴らしいです- DeepSpeed-ZeROでオプティマイザメモリをすでに保存していたため、8ビットオプティマイザを使用していませんでした）。

ブログのクレジット

以下の親切な方々に感謝の意を表します。良い質問をして記事の読みやすさを向上させてくれました（アルファベット順にリストアップ）：ブリトニー・ミュラー、ダウェ・キエラ、ジャレッド・カスパー、ジェフ・ラズリー、ジュリアン・ローネイ、レアンドロ・フォン・ヴェラ、オマール・サンセビエロ、ステファン・シュヴェーター、トーマス・ワン。

メインのグラフィックはチュンテ・リーが作成しました。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!

Was this article helpful?

93 out of 132 found this helpful

BLOOMトレーニングの技術背後