「Google は依然として最高です。」仕事を辞めて自分のビジネスを始めて 1 年経った後、大規模なモデルのトレーニングには多くの落とし穴があることに気づきました。

1 年以内に会社を設立し、資金を調達し、チップを購入し、Gemini pro/GPT 3.5 に追いつく LLM を構築するにはどうすればよいでしょうか。

インフラストラクチャの構築や、大規模な言語モデルやマルチモーダル モデルのトレーニングに興味を持つ人は多くいますが、実際に「ゼロから始める」というプロセスを経験する人はほとんどいません。一般的に、技術的な才能を確保することが前提条件であり、コアアルゴリズムを習得することが鍵であると信じていますが、実際には、エンジニアリングの実践で生じる課題は本当に頭痛の種です。

1年前、大規模モデルへの流行に乗って、Yi Tay氏は3年以上勤務したGoogleを辞め、Rekaという会社を共同設立し、大規模言語モデルに重点を置く主任科学者として活躍した。

Google 在籍中、Yi Tay は PaLM、UL2、Flan-U-PaLM、LaMDA/Bard、ViT-22B、PaLI、MUM など、数多くの有名な大規模言語モデルやマルチモーダル モデルに取り組みました。これほど豊富な経験があったにもかかわらず、彼はこれまで想像もできなかった困難に直面しました。より多くの起業家が落とし穴を避けられるように、Yi Tay 氏は自身が遭遇した落とし穴をブログ記事で共有しました。

「コンピューティングの不足と信頼性の低いコンピューティング プロバイダーにより、予想以上に困難が生じましたが、当社のテクノロジーの力で乗り切ることができました。最後に、いくつかの課題と学んだ教訓を明らかにするために、このブログ記事を書きました。この投稿が多くの人にとって興味深く、教育的なものとなることを願っています。」

記事が公開された後、多くのテクノロジー起業家によって議論され、転送されました。

アンドレイ・カルパシーも同意する。

成熟した企業には、クラスターを維持するための専任チームがあります。規模が大きくなるにつれて、クラスターはエンジニアリングの範囲から離れて生物学的なものになるため、「ハードウェアの健全性」に特化したチームが必要になります。

トレーニング実行の「ベビーシッター」は、大規模なモデルをトレーニングする際の日常的なストレスの多い作業です。実行のバイタルサイン（損失の急増、数値の問題、スループット、勾配基準、ポリシーエントロピーなど）を注意深く監視する必要があります。パフォーマンスが低下したり、停滞したりしたときは (頻繁に発生する可能性があります)、スタック トレースをすばやく調べて何が起こったかを確認する必要があります。この作業は迅速に行う必要があります。そうしないと、10,000 個の GPU がアイドル状態になる可能性があります。通常、それはこれまで見たことのない新しい、奇妙で恐ろしいバグなので、誰かが問題を見つけられるかどうかを確認するために助けを求める必要があります。

最悪のミスは午前4時に起こりました。通常、誰もそれを見ることはできないので、疑わしいと思われるノードをいくつか禁止して、再度稼働させようとすることができます。時には、その日に神々の祝福を受けなかったというだけの理由で実行が失敗することもあります。その場合は、起動コマンドに while True: ループを追加します。潜在的な問題は、GPU が過熱して時々間違った乗算演算を実行する、一部のルーターがダウンしてネットワーク ファイル システムの I/O が減少する、データ センター内の誰かが連絡なしのメンテナンス中に物理的に配線を切断するなど、多岐にわたります。あなたには決して分からない質問がいくつかあります。

明るい兆しを見つけた人々もいた。イー・テイ氏が言及した「荒野」とは「Google以外の企業」を意味するのだ。

インフラやハードウェアの観点から言えば、Google に匹敵するチームは多くありません。

それでは、ブログの内容を見てみましょう。

LLM時代のハードウェア抽選

モデルをトレーニングするための最初の要件は、コンピューティング能力へのアクセスです。これはかなり単純なことのように思えますが、最大の驚きは、コンピューティング プロバイダーの不安定性と、ソースに応じてクラスター、アクセラレーター、およびそれらの接続の品質に大きな違いがあることでした。

人々は、これは単にアクセラレータの選択に関する質問/議論（TPU 対 GPU など）であり、すべての GPU クラスターは同じであると考える傾向があります。私たちの経験では、これはすぐに間違っていることが証明されます。

さまざまなサービス プロバイダーをサンプリングしたところ、同じハードウェア (GPU (H100)) でも、ハードウェアの品質は大きく異なることがわかりました。ここでのハードウェアは、必ずしもチップやアクセラレータ自体ではなく、クラスターの全体的な品質を指すことに注意してください。全体的な感覚としては、宝くじを買うような感じです。

具体的には、複数のコンピューティング プロバイダーから、それぞれ数百から数千個のチップを備えた複数のクラスターをレンタルします。私たちは、問題のないクラスター（修正に数時間しかかからない小さな問題のみ）から、さまざまな理由で数時間ごとに障害が発生する完全に使用できないクラスターまで、さまざまなクラスターを見てきました。具体的には、ケーブルの問題 (N が不当に小さい)、GPU ハードウェア エラーなどにより、一部のクラスターでは N 時間ごとにノードが故障していました。さらに驚くべきことは、同じプロバイダーの各クラスターでも堅牢性が大きく異なる可能性があることです。

一方、他のクラスター ノードが大幅に安定している場合でも、I/O やファイル システムの問題が発生する可能性があり、チェックポイントの保存でもタイムアウトしたり、時間がかかりすぎてクラスターの使用率が低下する可能性があります。他のコンピューティング リソースの中には、実行するためにまったく異なるソフトウェア レイヤーを必要とするものもあり、独自のコードベースを持ち込むチームには適していません。実験や大規模なジョブを実行するには、追加の移行コストが必要になります。

完璧なものはありませんが、プロバイダーが提供するサービスの品質は異なることは確かです。

最もイライラすることは何ですか?特に、すべてが整った後では、どのようなハードウェアが提供されることになるのか、エクスペリエンスがどの程度強力でフォールト トレランスがあるのか​​を事前に知ることはほぼ不可能です。

さらに、サプライヤーが納期どおりに納品できず、機器の納品が数か月遅れ、ユーザーが数週間または数か月間他のソースから購入できなくなると、何が起こるかわかりません。プロバイダーによっては、チェックポイントを誤って削除してしまうこともあります。

異なるクラスターには異なるモデル フリップ使用率 (MFU) があることを述べましたか?運悪く、ノードの配線が不十分であったり、その他の問題を抱えているプロバイダーを見つけた場合、無駄にされるコンピューティングの量は無視できないものになります。システムのファイル システムが最適ではない場合、チーム メンバーがクラスター間で大量のデータを転送し始めると、トレーニング実行の MFU が低下します。

各サービスプロバイダーが提供するアフターサービスのレベルも異なります。それらは、丁寧なものから中途半端なものまで、「会話的な」定型的な返答からすべての問題をユーザーのせいにする回答まで多岐にわたります。

要約すると、私たちが試したすべてのクラスターには、独自のスタイル、問題点、および障害モードがあります。さらに、ほぼすべてのクラスターでは、さまざまな問題に対処するために独自のホットフィックス セットが必要です。それでも、フェイルセーフティは非常に重要であり、あらゆるクラスターに対して迅速なホットフィックス ソリューションを見つけることが重要であると認識しています。

過去数か月間、私たちは監視、効率的なチェックポイント、その他のさまざまな最適化など、すべてが利用可能であることを確認するための多くのツールを構築し、スケーラブルなデータ ストレージ用のカスタム ファイル システムもインストールしましたが、これは氷山の一角にすぎません。

これらのツールを組み合わせることで、MFU が大幅に改善されるとともに、ハードウェアの状態が悪い場合のダウンタイムも最小限に抑えられます。

GPU と TPU

私の会社では、ほとんどの時間を GPU を使用したモデルのトレーニングに費やしています。しかし、Reka に入社する前は、Google で大規模な言語モデルのトレーニングに TPU を使用していました。 CUDA と nccl は私にとって最も馴染みのないものです (Nvidia で働いていた同僚から、これが Nickel と発音されることを知りました)。

Google での TPU の経験と比較すると、GPU の故障率にはまったく驚きました。実際、大規模な実行でも TPU が何度も故障した記憶はありませんが、優れたインフラストラクチャと専用のハードウェア チームがあるために、このことに気付いていないだけなのかどうかはわかりません。実際、Google の UL2 20B モデルは、偶然に 1 か月間実行することでトレーニングされました。 GPU 分野であれば、最初の数日以内に確実に失敗していたでしょう。

そうは言っても、これは基盤となるシリコンよりも、アクセラレータを管理するハードウェア チームの能力に関係しているのではないかと思います。優れたハードウェア サポート (コンピューティング プロバイダーから) を受けることは非常に重要です。これは、彼らが本当に能力があるかどうかに大きく依存しており、「ハードウェア宝くじ」の概念を裏付けています。

GPU スペースが変な感じがします。 TPU ポッドでの分散トレーニングの最高水準と比較すると、マルチノード トレーニングは後付けのものです。 GPU 分野では、さまざまなプロバイダーがさまざまな方法で GPU を接続してマルチノード トレーニングを可能にしているように感じられ、場所によって実践方法が大きく異なります。

私はハードウェアの専門家ではありませんが、これが私の率直な感想です。

マルチクラスタ設定の苦労

私はキャリアのほとんどを Google インフラストラクチャに費やし、主に Borg、Xmanager、Colossus 上で実行してきました。そのため、異なるクラスターに新しい環境をセットアップする必要があるという概念は、私にとって非常に馴染みのないものでした。

多数のアクセラレータ プールが 1 つの場所に特別に設置されない限り、アクセラレータ プールのクラスターが複数存在することは、現代では避けられないと思われます。より具体的には、GPU の供給 (またはその欠如) によって、物事の性質が断片化されるこのクラスター購入モデルが自然に作成されます。大規模なモデルをトレーニングするにはテラバイト単位のデータも必要であり、データを移動するだけでも不便な場合があります。同時に、データのコピーは通常は簡単な作業ではなく、超大規模な場合にはデータのコピーコストも非常に高くなります。

明らかに、理想的な状況は、ジョブを異なるサーバーに具体的に送信する何らかのオーケストレーション レイヤーを備えることです。 AI に注力している大企業の多くは、一般的に研究者の生活の質を向上させるための何らかのインフラを整備していると思います。しかし、スタートアップが最初からこのような複雑で高度な ML トレーニング インフラストラクチャを構築するのは実際には不可能です。

私たちは、これらの問題を軽減し、世界クラスの実験インフラストラクチャのゴールドスタンダードに向けて前進し続けるために、いくつかの内部ワークフローを開発しました。 （この粗末な体制は、トップ企業や大企業以外ではほぼ標準だと聞きました）。

ワイルドコード

皆さんご存知のとおり、私がこれまでで最も気に入っているライブラリは T5X と Mesh Tensorflow ですが、これらにはいくつか欠点があります。

1) Google 以外ではそれほど多くのサポートがありません。

2) やや非推奨となっている。

3) 彼らは私たちのグループ内の非 xoogler に対して友好的ではありません。

最終的に、一般的で、安定していて、より人気のあるもの、つまり pytorch を選択しました。 PyTorch は、チームのほとんどの人（私以外）にとって、はるかに使いやすいです。

最初の数か月は、pip、git、docker などの「奇妙な」ものに戸惑いました。とはいえ、Google コードベースを外部で使用することがどれほど安定しているか、またはユーザーフレンドリーであるかは、100% 確信がありません。

率直に言って、外部コードベースの品質は、私が Google で慣れ親しんでいたものに比べて大幅に遅れていました。主な理由は、Google 内部のコードベースは ML の専門家自身 (Noam Shazeer、Barret Zoph、Adam Roberts、Hyung Won Chung など) によって作成される傾向があり、私が外部で試したコードベースと比較して優れていると感じるからです。

さらに、あるコードベースでモデルの並列性を変更するトランスフォーマーを記述する必要があるまで、モデルの並列性を変更する機能が自動的 (無料) ではないことを知りませんでした。これは私にとって間違いなく「WTF の瞬間」でした。

驚くべきことに、これらのコードベースでは、大規模なエンコーダー/デコーダーのトレーニングや、prefixLM トレーニングがほとんどサポートされていません。

原則は少なく、Yoloは多く

モデルを体系的に拡張するには、多くの場合、原則に従って小規模から始めて大規模に拡張し、複数の段階（10億→80億→640億→3000億など）で実験を実行し、成功したものを選んで拡張する必要があります。スタートアップでは、大規模なスイープを実行して、はるかに少ない計算量でハイパーパラメータをチェックします。結局、Yolo を複数回実行する必要がありましたが、幸運にも結果は良好でした。

最終的には、強力な 21B Reka Flash および 7B Edge モデル、そして今後発売される Max Core モデルを実現するために、小規模かつ短時間のアブレーション実行を数回行うだけで済みます。非常に限られた実行回数で信頼性の高いソリューションを見つけるのは困難であり、検索空間が非常に大きいため、一度に多くの変数を変更する必要があります。これを実現するには、人々は大手テクノロジー企業の体系的な性質を捨て去り、「YOLO」、つまり直感と本能に大きく頼らなければなりません。

ありがたいことに、私や私たちのチームの多くのメンバーは、ML キャリアを通じて十分な「直感」を培い、非常に短い時間で正しい答えを導き出すことができました。これまでの作業では非常に優れたモデルをトレーニングしてきましたが、トレーニング インフラストラクチャ、データ、新しいアイデアの組み込み、その他の環境の問題の違いにより、依然として結果に大きな違いが生じる可能性があります。そうは言っても、強力な事前確率は検索空間を大幅に削減するのに役立ちます。これが、非常に少ない試行、リソース、実験で非常に強力なモデルをトレーニングできる理由の 1 つであると考えられます。