re:Invent 2024: AmazonがSageMaker HyperPodでNova FMを構築した方法
はじめに
海外の様々な講演を日本語記事に書き起こすことで、隠れた良質な情報をもっと身近なものに。そんなコンセプトで進める本企画で今回取り上げるプレゼンテーションはこちら！
📖 AWS re:Invent 2024 - A close look at how Amazon built the Nova FMs using SageMaker HyperPod (AIM379)
この動画では、Amazon NovaモデルのローンチとSageMaker HyperPodについて解説しています。Foundation Modelの開発における主要な課題として、メモリの壁と計算の壁、そしてシステムの複雑さに起因するエントロピーの問題が取り上げられます。特に5万台以上のAcceleratorを使用する大規模クラスターでは1日に10-20回の障害が発生する現実があり、その対処方法としてBurn-inプロセスやGoodputの測定、迅速な障害復旧などの手法が紹介されます。これらの知見を活かしたSageMaker HyperPodは、自動障害検知や複数のオーケストレーションツール、事前最適化されたRecipesなどを提供し、Foundation Model開発における様々な課題を解決します。
※ 動画から自動生成した記事になります。誤字脱字や誤った内容が記載される可能性がありますので、正確な情報は動画本編をご覧ください。
※ 画像をクリックすると、動画中の該当シーンに遷移します。
re:Invent 2024関連の書き起こし記事については、こちらのSpreadsheet に情報をまとめています。合わせてご確認ください！
本編
Amazon Novaモデルのローンチと本セッションの概要
みなさん、こんにちは。Re:Inventを楽しんでいただけていることと思います。本日は、Amazon Novaモデルのローンチについて、大変エキサイティングなお知らせがあります。このセッションでは、これらのモデルがどのように開発されたのか、どのような課題に直面したのか、そこから得られた教訓は何か、そしてそれらの教訓をSageMaker HyperPodにどのように製品化したのかについてお話しします。これにより、Foundation Modelの構築やファインチューニングを行う際に、皆様もこれらの知見を活用することができます。
私はSageMaker Trainingチームを率いるRajneesh Singhです。本日は、同僚のAnand RathiとGordon McShaneが同席しています。AnandはAmazonのArtificial General Intelligence Group（通称AGI）のDirector of Engineeringで、GordonはAGIグループのPrincipal Engineerです。本日の参加者の皆様について、少しお伺いしたいと思います。Foundation Modelの構築やファインチューニングの経験がある方は手を挙げていただけますか？ - およそ25%ですね。では、Foundation Modelがどのように構築されているのか知りたいと思っている方は手を挙げてください - 残りの皆様ですね。素晴らしいです。今日のセッションから、皆様が実践で活用できる知見を得ていただけることを願っています。
それでは、早速始めましょう。このセッションでは、Foundation Model開発における課題についてお話しします。AnandがAmazon Nova Foundation Modelsの概要について説明し、GordonがAGIチームがこれらの課題にどのように対処したかについて説明します。そして私が、Amazon SageMaker HyperPodがFoundation Model開発のこれらの課題をどのように解決するのかについてお話しします。
Foundation Model開発の課題とAmazonのAI活用事例
まず、Foundation Model開発における課題を大まかに見ていきましょう。Foundation Modelを構築またはファインチューニングする場合、まず膨大な量のデータを調達する必要があります。データを調達した後、それを複数のチャンクに分割し、クラスター内の数百から数千のアクセラレーターにそれらのチャンクをロードする必要があります。さらに、複数のトレーニング手法を試し、トレーニングスクリプトを作成し、数ヶ月かけてスクリプトを改良し、チェックポイント戦略を立て、ハードウェア障害に対処する必要があります。これは非常に大変な作業です。
ここで、AGIがこれらのFoundation Modelをどのように構築し、これらの課題に対処したのかについて、Anandにお話しいただきたいと思います。Anand、お願いします。ありがとうございます、Rajneesh。皆様、こんにちは。私はAnand Rathiで、AGIチームのDirector of Software Developmentを務めています。私はAmazonに20年以上在籍しており、現在の役割では、高性能コンピューティングインフラストラクチャーとMLツーリングの開発を率いるチームを指揮しています。これらは、AGIの科学者たちがNovaモデルを構築・トレーニングする際に使用するものです。
まずは少し歴史を振り返ってみましょう。Amazonの創業以来、当社のミッションは常にお客様の生活をより良いものにすることでした。AIとMLシステムは、このビジョンとミッションを実現する上で重要な役割を果たしてきました。皆さんもAmazonのパーソナライズされたレコメンデーションを経験されたことがあるはずです。「商品Xを購入したお客様は、こちらの商品にも興味を持っています」というのは、大規模なML基盤のレコメンデーションシステムの最も初期の導入例の一つで、お客様が思いもよらなかった新しい商品を発見できるようになりました。
Alexaは最先端のパーソナルデジタルアシスタントですが、その裏側では30の異なるMLシステムが連携して、毎週数十億件にも及ぶお客様からのクエリに応答しています。これには天気予報やその他の情報の問い合わせ、音楽やビデオコンテンツの発見、スマートホームデバイスの制御など、実に様々な機能が含まれています。 Prime Airのドローンは、複雑な地形をナビゲートするためにコンピュータビジョンモデルを使用しており、60分以内での商品配送を目指したパイロットプログラムを開始しています。
これは非常にエキサイティングで印象的な取り組みです。当社のフルフィルメントセンターでは、何十万台ものロボットと関連するMLシステムが連携して、 毎日何百万個もの商品をお客様にお届けしています。そしてこれは、会社全体で見られるAI技術活用の氷山の一角に過ぎません。私たちは、お客様体験を継続的に向上させるために、AI基盤の技術の採用を積極的に進めています。
Amazon Nova Foundation Modelsの概要と活用事例
これらの何十年にも及ぶ経験と専門知識を活かし、Amazon Novaファミリーのファンデーションモデルを構築し、お客様と共有することにしました。Amazon Novaファミリーの主力モデルは、Nova Micro、Nova Lite、Nova Proです。Microはテキストのみの入出力に対応した超高速なファンデーションモデルです。より高度なニーズには、LiteまたはProへステップアップすることができます。これらはテキストベースの出力を維持しながら、入力モダリティとしてテキストに加えて画像や動画もサポートしています。
これら3つのモデルに加えて、Nova CanvasとReelというテキストから画像、テキストから動画を生成するファンデーションモデルもあります。お客様は簡単なテキストプロンプトで画像や動画を生成でき、さらにカメラパンニングなどの高度なコントロールも可能で、モデルの出力をニーズや要望に合わせて調整できます。来年第1四半期には、最も高度な推論能力を持つマルチモーダルモデルである、Amazon Nova Premierをリリースする予定です。
社内では、さまざまなチームやビジネス部門でこれらのモデルの採用が着実に進んでいます。 あらゆる規模のブランドが、販売する商品の簡単な説明を入力するだけで、数回のクリックで顧客の心に響く広告キャンペーン用のクリエイティブを生成できるようになりました。これにより、すぐに広告キャンペーンを開始することができます。 同様に、Prime Videoではエピソードやシーズンのダイジェスト動画の自動生成にNovaモデルを活用しています。新しいシーズンを視聴し始める際に短いダイジェスト動画が表示されますが、従来はその制作に数週間の手作業が必要でした。しかし、これらのモデルを使用することで、動画を取り込み、重要なプロットポイントやストーリーの展開を特定し、ダイジェスト動画の推奨内容を生成できるため、手作業の時間を数週間から数日に短縮することができます。
Developer toolsは、これらのモデルを展開する上で非常に有望な分野です。コーディングアシスタントとしても、運用アシスタントとしても、幅広い活動を行うことができます。エンジニアが運用上の問題に対応するために呼び出されたとき、ダッシュボードやデプロイメントから最新のメトリクスなどのデータを自動的に収集することができます。エンジニアが調査を開始する時点で、すぐに利用できる要約が用意されており、トラブルシューティングの時間を短縮できます。さらに、ドキュメントの作成やユニットテストの生成といった関連タスクも実行できるため、開発者の生産性を大幅に向上させることができます。
他の運用チームもこれらのモデル群から大きな成果を得ています。例えば、AWS Field Engineeringは様々なAWSカスタマーをサポートしており、Genieというシステムを使用しています。現在、Novaの機能により、Genieはお客様の複雑なシステムやアーキテクチャ図を取り込み、AWSのベストプラクティスやドキュメントに基づいて最適化の機会を特定し、Field Engineerがより迅速にお客様をサポートできるようになっています。私たちは、お客様がこれらのモデルで何を実現できるのかを見るのが本当に楽しみです。
Foundation Modelの構築プロセスと直面する課題
これらのモデルの構築には多大な労力が必要なので、その過程について皆様に概要をお話ししたいと思います。Rajneeshが言及したように、まず大量のデータを収集することから始まります。そして、言語や形式の両面で、データの多様性を確保することが重要です。データが利用可能になったら、それを処理する必要があります。この段階で、関連性のない内容、有害なコンテンツ、バイアス、重複コンテンツをフィルタリングします。その後、正規化やトークン化の技術を適用し、利用可能な形でデータをパッケージ化します。これと並行して、サイエンスチームは、構築したいモデルの種類を特定するために、さまざまなモデルアーキテクチャやハイパーパラメータのチューニングに関する実験を行います。これらの2つのステップが完了すると、大規模なトレーニングを開始する準備が整います。
トレーニングプロセスは、Pre-trainingフェーズとPost-training（Fine-tuning）フェーズの2つの部分で構成されています。Pre-trainingは最も計算負荷の高い部分で、モデルアーキテクチャと様々なHyperparameterを使用します。これらのトレーニングレシピを使って、データを取り込み、大規模なコンピュートクラスターを起動します。Amazon Nova Premierのようなモデルの場合、クラスターは数万から約10万のAcceleratorにまで及び、期間は数ヶ月に及ぶことがあります。システムメトリクスを監視して、異常や遅延、性能低下の兆候がないかを確認する必要があります。同時に、サイエンスチームはモデルの進捗状況を示すLoss曲線と、正しい軌道に乗っているかどうかを綿密に監視します。もし問題があれば、実行を一時停止し、データの組み合わせやパラメータを調整してから、その時点から再開します。
十分な量のトレーニングトークンが消費されると、モデルはFine-tuningステージに引き継がれます。Pre-trainingステージから出てくるのは、高度に一般化されたモデルです。Fine-tuningによってモデルをより特殊化し、指示に従う方法を学習させたり、会話AIの機能に特化させたりします。このステップが完了すると、高性能ではあるものの実行コストの高いTeacherモデルが得られます。 これらのモデルの推論コストを持続可能なものにするために、CompressionやDistillationなどの技術が適用され、同等の性能を持つより小さなモデルにパッケージ化されます。
社内外の様々なベンチマークでモデルのパフォーマンスに満足できると、 Bedrockなどのシステムを通じて提供される準備が整います。一つのサイズですべてに対応できるわけではないという理解に基づき、異なるサイズと機能を持つモデルのスイートを提供し、お客様が最適なパフォーマンスを得られるものを選択できるようにしています。時には最高のモデルでもニーズを満たせない場合があるため、お客様が独自のデータを持ち込んでこれらのモデルをカスタムFine-tuningし、特定のユースケースにさらに特化させることも可能です。
AGIに向けた競争が進む中、Scaling Ladderの特性により、モデルは大規模化の一途をたどっています。モデルサイズは2年ごとに4〜10倍の世代間ジャンプを見せています。ほんの数年前まで40億パラメータが最先端とされていましたが、現在では1兆パラメータの領域に達しています。しかし、Acceleratorを見てみると、メモリ容量は単に2倍になっているだけで、より大きなモデルのメモリ要求に追いつけないというメモリの壁に直面しています。2つ目の課題は計算の制約です。 2022年まで、モデルのトレーニングに必要なトークン数の決定は非常にヒューリスティックなプロセスでした。しかし、あるペーパーで、モデルパラメータ1つにつき20のトレーニングトークンが最適なトレーニングに必要であることが証明されました。それ以下のトークン数では、モデルは十分にトレーニングされていないことになります。
彼らは非常に興味深い比較を行い、700億パラメータのChinchillaモデルが1,750億パラメータのGPT-3を上回るパフォーマンスを示しました。この理論に従えば、700億パラメータのモデルを構築するには1.4兆のトークンが必要になります。これを将来に投影すると、XとY軸の両方が対数スケールである組み合わせ的な爆発が見えてきます。計算需要は指数関数的に増加し続けており、これだけではありません。
Chinchillaは、モデルパラメータに対するトレーニングトークンの比率として20対1が最適だと述べています。 しかし、Llamaの論文では、その比率を超えてもモデルは継続的に改善され、より賢くなることが証明されました。 現在のトレンドでは、モデルパラメータの40-60倍までトレーニングを行っています。8ビリオンパラメータのモデルでさえ、必要なデータと計算リソースがあれば、6-7兆トークンまでトレーニングすることが可能です。明らかにハードウェアの進化を待っているわけにはいかないので、私たちは大規模なクラスターを展開してこれらの課題に対処しています。より低い精度でのトレーニングやMOEによって得られるSparsityなど、他の科学的進歩も役立っていますが、トレーニングクラスターは避けられません。
大規模分散学習における技術的課題とその対処法
これまでに直面した課題とその対処方法について共有するために、Gordonに話を譲りたいと思います。ありがとう、Anand。皆さん、お越しいただきありがとうございます。私はAGI組織のPrincipal Engineerです。Anandの下で働いており、インフラストラクチャに焦点を当てて取り組んでいます。本日は、特に私たちが直面したインフラストラクチャの課題についてお話しさせていただきます。
具体的な話に入りましょう。 まず、私たちが持っている物理インフラについて説明し、その後で問題点に移ります。最下層には、コンピュートとストレージがあります。ここではAmazon EC2のアクセラレーテッドインスタンスを使用しており、大規模なフリートを運用しています。P4D、P5、そしてAndyが先ほど発表したTrainium 1とTrainium 2など、様々な組み合わせを使用しています。また、これらの高性能インスタンス間の相互接続にEFAを使用しています。これらのパワフルなインスタンスはIO boundなので、EFAは非常に重要です。Elastic Fabric Adapterはマイクロ秒レベルのレイテンシーと、毎秒数百ギガビット、50ギガバイトの帯域幅を提供します。
右側には、ブロックストアとしてAmazon FSx for Lustreがあります。これは非常に高性能なHPCクラスのネットワークアタッチストレージで、様々な目的で使用していますが、主にデータのシャッフルやリパッキングをオンザフライで行う必要がある場合に役立ちます。これには通常、大量のランダムアクセスが必要になるためです。そして、すべてのバックアップにAmazon S3を使用しています。これは私たちのオブジェクトストアであり、セキュアで高い信頼性があるため採用しています。長年にわたり、これをほぼ無限にスケールアウトする方法を学んできました。S3チームは新しいExpress Zoneで、さらにローカルで高性能な低レイテンシーの容量を提供する機能を追加し続けています。チェックポイントをそこに保存し、多くのマルチモーダルデータをS3からストリーミングしています。
これらの上に、Amazon SageMaker HyperPod製品ですべてを統合しています。下層でこれらのインフラストラクチャをすべて管理することは非常に時間がかかります。これらのスタック層にはそれぞれ異なるドライバー - CUDAバージョン、Neuronバージョン、EFAバージョン、これらすべてのコンポーネント、カーネルバージョン、NCCLバージョンなど - があり、すべてが適切な安定したパフォーマンスを発揮する構成である必要があります。HyperPodは私たちのコントロールプレーンであり、これらのクラスターを簡単にデプロイし、維持し、すべてのインスタンスをオンザフライでアップグレードし、さらに監視することができます。
今年この製品で特に印象的だったのは、SageMaker HyperPodチームが本当に理解してくれていたということです。私たちが必要とする機能に関して、今は選択肢があることが重要なのです。オープンソースコミュニティのスタックは、私たちがテストして試す必要のある多くの選択肢で活気づいています。HyperPodが実現したことの1つは、使用したいスケジューラーを自由に選べるようにしたことです。社内では、特にKubernetesとAmazon EKSに大きく依存しています。ここで少し、なぜKubernetesを使用しているのかについてお話ししたいと思います。
実は、ワークロードの多様性とホストタイプの異種性により、異なる設定間を素早く切り替える必要があるのです。Anandが言っていたように、ある日は蒸留を行い、次は事前学習を行い、さらに強化学習のような事後学習活動を行う必要があります。SageMaker HyperPod上のAmazon EKSは、このような全体的な設定を変更する際の俊敏性を本当に提供してくれています。
私たちはAWS上のApache AirflowのマネージドバージョンであるAmazon MWAAを、MLflowと組み合わせて使用することで、優れた実験の再現性を実現していることをお伝えしたいと思います。私たちの研究者たちはAirflowを愛用しています。なぜなら、トレーニングスクリプトとコードベースをPythonで定義でき、ハイパーパラメータの最適化やモデルトレーニングを行う際にもPythonを活用できるからです。MLflowはすべての設定やモデルメトリクス、システムメトリクスを記録するので、パフォーマンスのボトルネックやモデルのパフォーマンスの問題を詳しく調査することができます。右端には、マネージドPrometheus、マネージドGrafana、Cloud Watchがあり、これらに大量のメトリクスを送信しています。私たちは、ジョブレベル、推論レベル、クラスタースケジューリングレベルでパフォーマンスの遅延がどこで発生しているかを理解するために、これらのクラスターを徹底的にモニタリングする必要があるのです。
今日のトークで皆さんに理解していただきたい重要なポイントは、Foundation Modelのトレーニングにおける課題についてです。ここにいる方々の中には、すでに自分でモデルのファインチューニングを行っている方もいらっしゃると思います。これらの問題は基本的に同じです - 見た目は全く同じで、ただ規模と大きさが格段に拡大しているだけなのです。簡単な例えを使ってみましょう：車で移動する際に、H100をたくさん車の後部に積んで州間を走行するようなものです。ガス切れやタイヤのパンク、あるいは油圧システムやエンジン部品の交換のために車をディーラーまで押して行かなければならないような問題に遭遇するかもしれません。
しかし私たちにとって、これは実際には5万個ものアクセラレーターをジャンボジェットに積んで大西洋を横断するようなものなのです。ここで必要なのは、飛行機を着陸させることなくタイヤを交換し、給油できることです - どこかの土地に着陸するわけにはいきません。この飛行機を空中に保ち続けなければならないのです。
問題点と解決策について話しましょう。Anandが言及したように、いくつかの壁があります - メモリの壁と計算の壁です。これらの壁がどれほど高いのかを説明し、並列化戦略を使ってどのようにそれらを乗り越えられるかを検討します。これらの壁を乗り越えることは、ある問題の解決策ではありますが、それがまた新たな問題を生み出すことにもなります。そして、エントロピーの問題もあります。つまり、システムにコンポーネントを追加すればするほど、確実に苦労や問題が発生するということです。
まずメモリの壁から始めましょう。幸いなことに、これは線形的な問題です。モデルにパラメータを追加するたびに、それに応じたメモリ増加は一定です。具体的に見ていきましょう。 おそらく5年前に最先端だった40億パラメータのモデルを例に取ります。 Single PrecisionやBF16のモデルサイズに必要なバイト数を計算すると、重み、Activation、Optimizer stateを含めて約72ギガバイトのメモリが必要になります。
エントロピーと障害への対応：AGIチームの経験から
もちろん、これは家庭用ノートPCでは実現できない規模ですが、現代のハードウェアでは対応可能です。H100は現在80ギガバイトのRAMを搭載しています。P4dもかなりの量のメモリを搭載していました。ですので、おそらく収まるでしょう - Chromeのタブをいくつか閉じる必要があるかもしれませんが、なんとかなるでしょう。 では、現在のFoundation modelのサイズについて話しましょう。現在は4,000億パラメータです。この数字は多少恣意的に選んでいますが、現在の業界の状況をおおよそ反映しています。 18倍すると、7.2テラバイトになります。明らかに、これはもはや単一のH100には収まりません。実際には、約100台のH100か12台のP5が必要になります。悪くはありませんが、誰もが12台のP5を手元に用意できるわけではありません。これには波及効果もあります - モデルの重みとCheckpointを保存するのに7.2テラバイトが必要になります。
これらのCheckpointは可能な限り頻繁に、数秒単位で保存する必要があります。そうなると、ペタバイト規模のデータをストレージに保存することになり、それをどうやって高速かつ効率的に処理するかを考える必要があります。
次に計算の壁に移りましょう。こちらの状況はどうでしょうか？これは本日の講演のテーマになりますが、パラメータ数が増加するにつれて、必要な計算量は二次関数的、あるいはそれ以上に増加します。最適なトレーニングに必要な計算量は継続的に増加しています。ここでも同じような思考実験をしてみましょう。 40億パラメータのモデルを取り上げ、1,600億のTokenをそこに通すとします。これには約6回の浮動小数点演算が必要です。すると、 約3,840エクサフロップスになります。これは、この時点でもかなりの計算量が必要とされることを示しています。
先ほどの12台のP5で400億パラメータのモデルをロードしましたが、今度は40億パラメータのモデルを実際にトレーニングしてみましょう。すべてが完璧に進めば約1日でできることがわかりました。これは悪くない成果です。では、これを4,000億まで引き上げてみましょう。すると、16兆という数字になり、38百万エクサフロップスという計算量になります。同じ12台のP5で今これを行おうとすると、私の子供たちが大学に入学するころ...つまり約20年かかることになります。これは明らかに大きな問題であり、大きな課題です。
では、この問題をどのように解決するのでしょうか？実際には、かなり高度な並列化戦略を導入する必要があります。まずは、1台のGPUやNPUにすべてが収まる理想的な状態から始めましょう。最初に直面する課題は、このモデルの中の非常に計算コストの高いレイヤーに対処する必要があることです。特に、皆さんもご存知かもしれないAttentionレイヤー - Transformerモデルには、この特殊なAttentionレイヤーがあります。これらのレイヤーこそが、これらのモデルの二次的な性質を生み出している要因なのです。
この並列化が必要になります。モデルの幅が、必要な並列性を実際に示しています。これらのモデルにおける幅は、基本的にContext長に関連しています。皆さんがこれらのモデルを使用する際に馴染みのある、Bedrockウィンドウで送信ボタンをクリックする前にどれだけのContextを入力できるか、それが対処すべき幅を生み出しているのです。そこで、Tensor Parallelismと呼ばれる手法を使用します。個々のレイヤーを水平方向にスライスし、これらの計算を分散させます。この行列乗算を行った後、これらの結果をすべて集約する必要があります。これは実際にはネットワーク上でかなりの通信が発生し、特定のネットワークリンクに依存する必要があります。先ほど触れたNeuron Linkという非常に高速なインターコネクトがあり、NVIDIAも同様のNVLink技術を持っています。この種の並列処理を効率的に行えるのは、数百ギガバイト/秒のリンクを使用する場合のみです。
幅の問題が解決できたので、次は深さの問題です。これらのモデルは深くなればなるほど、より賢くなる傾向があります。実際、モデルが長く、より深くなるほど、タスクをより上手く処理できるようになります。そこで、Pipeline Parallelismと呼ばれるものを使用します。これは、レイヤーのシーケンスをステージ間で分割し、各ステージで計算を行い、次のステージに渡し、前のステージで次のサンプルの計算を続けるという方法です。これはかなり効率的な方法ですが、明確な非効率性もあります。上部の第3ステージがボトルネックになると、他のすべてのステージがそれを待つことになります。Pipeline Bubblesという用語を検索してみると良いでしょう。これがまさに、これらのレイヤーを待つことで発生する現象です。
ここまでで良い位置にいます - 先ほど話した12台のP5にすべてをロードできる状態になりましたが、これらの並列化をすべて行った後でも、トレーニング時間は約20年かかってしまいます。そこで、Data Parallelismという第3の並列化を導入する必要があります。これは実際にはかなりシンプルで、ほぼ水平スケーリングと言えます。GPUとNPUの2次元マトリックスを取り、モデルを複製します。1,000個のモデルを同時にトレーニングしているかのように、個々のデータシャードを渡します。そして、しばらくしてから全てのモデルを照合することを決め、すべての値を収集し、平均を取り、それらを戻します。
これを継続的に行うことができ、全体の同期コストが最大に達して収益逓減のような状態になるまで続けることができます。では、これらの問題の多くを解決できているように見えるのに、なぜ分散学習はこれほど難しいのでしょうか？
簡単な並列処理の問題を見てみましょう。過去20年間で、Googleが導入したMapReduceのようなもの（現在はApache Sparkとして形式化されています）を使って、人々がこの種の当たり前な並列処理を常時行っているのをご存知かと思います。これには3つの重要な特徴があります。通常、ステートレスであり、データの一部を処理するプロセスやワーカーは、新しいデータを計算する際に過去の処理を覚えておく必要がありません。また、ほとんど独立しており、ネットワークを介して相互に通信する必要がありません。自分の仕事をして、パイプラインに渡して、作業を続けることができます。さらに、一方向性があります - 入力は一方から、出力は他方からです。
分散学習の複雑性と障害対策の重要性
では、分散学習ではどうでしょうか？ 全く良くありません。高度にステートフルで、このグラフの各ノードは重要な情報、つまり私たちが成果物として必死に捕捉しようとしているモデルの重みを保持しています。また、高度に相互依存的です - Tensor ParallelismやData Parallelismで学んだように、学習を進めるには定期的な同期が必要です。学習自体に詳しい方なら、モデルの推論を行いLossを計算するForward Passがあり、その後、新しいLossを反映してすべての値を更新するためにBackward Passを行うことをご存知でしょう。
結局のところ、私たちに大きな課題をもたらすのはエントロピーです。素人的な定義を使うと、システムにコンポーネントを追加すると、故障点が増えるという現象です。システム全体が正しく動作することに依存している場合、システム全体が故障する統計的確率は増加します。私たちは数万台 - 5万、6万、7万、8万台のアクセラレーターを、CPUとメモリを搭載した数万台のホストと組み合わせています。そして、休憩や休暇なしで3〜4ヶ月間、可能な限り酷使して動かし続けます。
故障は単なる現実です。この規模では、1日に10〜20回の故障が発生するのが現実なのです。これはAmazon内部だけの話ではなく、文献を見ても他社も同様の状況を報告しています。私たちは故障率を下げる方法ではなく、それにどう対処するかを学ばなければなりません。これらすべてにどのように対処しているかを説明する前に、私が夜も眠れなくなるような問題、つまり実際に直面している問題をお話ししましょう。
このような障害が発生した場合の最良のシナリオは、GPUやNPUが突然消失してしまうというケースです。ドライバーがクラッシュし、デバイスが存在しなくなり、機能しなくなります。通常、PCIバスからの切断やNVIDIAのXIDエラーといった暗号めいたメッセージが表示されます。しかし、これは実はありがたいケースで、カーネルログメッセージに「私はもうダメです、他の誰かにこの作業を任せてください」というメッセージが表示されるからです。
次のケースは、それほど良くはありませんが最悪というわけでもなく、宇宙線のような外部要因に関連するものです。これらのフリーな粒子が実際にRAMに衝突してビットを反転させるという課題に直面しています。最近のハードウェアはECCメモリのような対策を講じていますが、これにより問題が発生したメモリ領域を再マッピングすることができます。しかし基本的には、これらのホストを再起動して対処する必要があります。修理が必要な2万台ものホストがある場合、これらの宇宙線の影響は徐々に大きな問題となってきます。
最後のクラス、これが私を本当に悩ませる種類の障害です。ハードウェアが何も報告してこない場合です。仕様通りに動作しなくなっているのに、それを示す信号がありません。業界ではSilent Data Corruption（SDC）と呼んでいます。これは文字通り、計算が正しく機能しなくなった状態で、GPUのプロセッサーでの乗算が微妙にずれていき、その結果、モデルが少しずつ破損していきます。Loss曲線やGrad normなど、すべてが狂い始めますが、これが実際に起きていることに気付くまでに何時間もかかることがあります。私が呼ばれた時には、モデルのトレーニングを停止し、2万から4万個のアクセラレーターの中から、どのホストが微妙におかしな動作をしているのかを特定しなければなりません。まさに干し草の山の中から針を探すようなものです。
ここで重要なのは、システム的なエラーメッセージのある障害に対処する方法を学ぶ必要があると同時に、症状しか分からず簡単には診断できない二次的な種類のエラーを検出できるようになる必要があるということです。では、このようなエントロピーや障害とどのように付き合っていけばよいのでしょうか？私たちは早い段階で、カオスを受け入れ、それを前提とした設計をする必要があることを学びました。早期の計画が重要で、私たちが行っているのはBurn-inと呼ばれるプロセスです。この用語は製造業から借用したもので、厳密には異なりますが、考え方は同じです。非常に特殊な決定論的なワークロードでロードテストを実行し、ハードウェアの製造上の影響を早期に特定しようとします。これにより、正常に動作しないハードウェアがクラスターに入り込むことを防いでいます。
その後、受動的および能動的な監視を2つの方法で開始します。カーネルログ、ドライバーログ、およびこれらのアクセラレーターから出力される他のメトリクスを監視します。また、特定のアクセラレーターが他のアクセラレーターと非常に遅い速度で通信している理由など、二次的な信号も監視します。以前にベースラインを取得しており、10倍速くなるはずだと分かっているため、そのアクセラレーターが故障しかけているか、何らかの問題が発生しそうだということが分かります。また、時間を節約するために効率的かつ非同期的に保存することも学びました。Amazonには冗長な容量を持つという贅沢があるため、ブレークポイントで待機している予備のホストを用意しておくことができます。エラーが発生した場合、すぐに新しいホストに切り替えられる態勢が整っています。
私たちはモデルからホストまで、スタックのあらゆる層から膨大な量のメトリクスを測定し、収集しています。ここで重要なのは、これだけのデータを収集した後、それを意味のある形で可視化する必要があるということです。データ量が非常に大きいため、実際のモデル、クラスター、データセンターの形を反映した意味のあるダッシュボードに変換しない限り、効果的な分析はできません。これは現在、AWSでML Topology APIを使用している人なら誰でも実現可能です。このように意味のある可視化を行うことで、クラスターの特定の領域に問題があるかどうかを素早く判断できます。
3番目に重要なのは、KPIを把握し、これらすべてのメトリクスをGoodputのような高レベルのKPIに落とし込むことです。Goodputは、これらのモデルが稼働し、実際に価値のある処理を行い、全体の出力に貢献している時間の割合を示すために測定しています。チェックポイントの保存、インフラの再起動、モデルのメモリへの再ロード、データの再読み込みなどはカウントされません。私たちは常に95-98%のGoodputを目指しています。
最後に、私たちは素早く失敗し、素早く回復する必要があります。何か問題が見つかった場合、待つことはせず、すぐにクラッシュさせます。私たちが優先するのは回復であり、絶対に必要なものだけを停止します。トレーニングプロセスで40-50,000個のAcceleratorを使用している場合、それぞれに個別のプロセスが関連付けられています。もしそれらすべてのプロセスとコンテナを停止させると、再起動に10分ほどかかってしまいます。しかし、私たちはそうせず、個々のプロセス自体だけを再起動します。
また、モデルデータのインデックスなどをキャッシュすることで、クラッシュから数秒で起動できるようにしています。そして先ほども触れましたが、このスケールではチェックポイントの頻度を最適化することが重要です。もし1時間や2時間ごとにしかチェックポイントを取らない場合、10分から20分の間隔でクラッシュが発生するたびに2時間前に戻る必要があります。回復にそれほど時間がかかると、実際に時間を失うことになってしまいます。
これらすべての素晴らしい点は、私の苦労があなたの利益になるということです。私はRajneesh と彼のチームと多くの時間を費やし、これらの機能について協力してきました。そしてそれらはすべてHyperPodに直接組み込まれています。では、彼が戻ってきて、これらをプロダクトに組み込むために行った素晴らしい取り組みについて説明してくれます。ありがとう、Gordon。GordonからAGIチームが直面した課題とその解決方法について説明がありました。良いニュースは、これらの学びをすべてSageMaker HyperPodに製品化し、皆さんに提供できるということです。そのため、皆さんが自分でこれらの課題や問題に直面する必要はありません。
SageMaker HyperPodの機能と最新リリースの紹介
SageMaker HyperPodを使用することで、エントロピーの問題、メモリウォール、あるいはコンピュートウォールに対処することができます。SageMaker HyperPodは耐障害性の高い環境で、詳細なヘルスチェックが組み込まれた自己修復型クラスターを提供します。Data Parallel、Tensor Parallel、Pipeline Parallelなど、どのような並列化手法を使用する場合でも、SageMaker HyperPodには最適化されたライブラリが製品の一部として組み込まれ、統合されています。データサイエンティストがジョブの投入や管理のためのユーザーインターフェースを使用したい場合は、SageMaker Studioというツールでそれらのインターフェースも利用できます。
データサイエンティストがコマンドラインインターフェースを使用したい場合は、HyperPod CLIも利用可能です。MLOpsエンジニアなどのオペレーターがContainer Insightsを必要とする場合も、SageMaker HyperPodを通じて利用できます。それでは、SageMaker HyperPodの主要な要素を見ていきましょう。 SageMaker HyperPodは、EFAを備えたEC2 Ultraクラスターと高性能ストレージを提供します。 インフラストラクチャとソフトウェアの両方の要素が、HyperPodを通じて直接提供されるAmazonインフラストラクチャ向けに最適化されています。
これらのレイヤーを見ると、障害はほぼどこでも発生する可能性があります。インフラストラクチャレイヤーでも、ソフトウェアレイヤーでも発生する可能性があります。HyperPodがどのように対処するか見ていきましょう。HyperPodは障害のあるノードを自動的に置き換える機能を提供します。 トレーニングジョブを開始すると、チェックポイントも作成されます。これらのチェックポイントが作成される際、モデルの状態やトレーニングプロセスの状態が保存されます。ノード障害が発生した場合、HyperPodは自動的にノード障害を検出し、正常なノードに置き換えます。正常なノードに置き換えた後、前回の安全なチェックポイントに戻り、トレーニングプロセスを再開します。
これらすべてが手動介入なしで実行できます。そのため、このような障害に自分で対処する時間を費やす必要がありません。お客様からは、この機能によってFoundation Modelのトレーニング時間を大幅に短縮し、トレーニングコストも大幅に削減できているとの声をいただいています。自動的な障害ノードの置き換えを上書きして手動で行いたい場合も、そのオプションは利用可能です。 障害対応の後は、GPUであれ、クラスター自体の利用率であれ、この高価なインフラストラクチャの利用率を向上させる方法も確認したいでしょう。HyperPodは、TrainiumとGPUの両方についてアクセラレーターレベルの可観測性を提供します。
この可観測性を使用することで、トレーニングスクリプトを改善し、利用率を最大化することができます。同様に、クラスター側では、複数のチームが同じクラスターを使用し、複数のトレーニングジョブを実行できます。これらのトレーニングジョブに優先順位を付けて配分することで、より高い利用率を実現できます。これらのツールは非常に便利で、 Gordonが説明していたように、AGIチームもこれらのツールに依存しているほどです。
HyperPodは大きな柔軟性を提供します。オーケストレーションを例に取ると、HyperPodはオーケストレーターとしてAmazon EKSを提供し、 さらにSlurmもオーケストレーターとして提供しています。また、様々な種類のジョブ投入ツール、観測機能、トレーニングおよびチューニングのインターフェースを提供しています。 さらに、先ほど説明したように、PyTorch、TensorFlow、NEMO、JAXなど、様々なトレーニングライブラリやフレームワークを提供しています。 また、HyperPod向けに最適化された多くのオープンソースの分散トレーニング戦略や、複数のMPSツールも提供しています。
ソフトウェアとドライバーに関しては、すぐに使えるプリビルドイメージが用意されており、多くのデバイスドライバーも利用可能です。 ツールキットについて、そしてハードウェア構成を見てみると、HyperPodは 様々な選択肢を提供しています。H100やA100を搭載したP4またはP5インスタンスのAmazon EC2インスタンス、複数のストレージオプション、EFAによるネットワーキングなどを利用できます。これらは、お客様の好みや要件に応じて選択できる様々なオプションです。
また、Foundation Modelの構築にSageMaker HyperPodを利用している主要なお客様についてもご紹介したいと思います。Luma AIやPerplexityなどの多くの有力スタートアップ企業がHyperPodを使用してFoundation Modelの取り組みを拡大しており、SalesforceやThomson Reutersなどの企業もHyperPodを使用してFoundation Modelを構築しています。
最近のリリースについて見ていきましょう。re:Inventでは複数の新機能をリリースしましたが、特に 私が注目している2つのリリースについてお話ししたいと思います。1つ目は、今朝リリースしたHyperPod Task Governance Toolsです。これは、複数のチーム間でコンピュートリソースを分散させ、トレーニングジョブの優先順位を定義できる機能です。例えば、同じクラスターを推論とトレーニングのワークロードに使用している場合、日中は推論トラフィックが多いため、そのワークロードにより多くのコンピュートリソースを割り当てたいと考えます。夜間に推論のトラフィックが減少した際には、そのコンピュートリソースをトレーニングジョブに割り当てることができます。HyperPodがこれを管理し、これらのパラメータを調整する機能を提供します。
次にお話ししたいのはHyperPod Recipesです。 HyperPod RecipesはAWSインフラストラクチャー向けに最適化された、事前にベンチマーク済みのレシピです。LlamaやMistralなどのオープンソースモデルを使用する場合、これらのモデルは特定のアーキテクチャに従っており、Foundation Modelのトレーニングや微調整のために調整が必要です。通常、データサイエンティストはこれらのパラメータの微調整に数週間、場合によっては数ヶ月を費やします。適切なバッチサイズやオプティマイザーの設定の決定などがその例です。HyperPod Recipeの機能により、オープンソースモデル向けに約30種類のレシピを提供しており、これらをテンプレートとして直接使用することで、最初から最高クラスのパフォーマンスを実現できます。
セッションの終わりに、重要なポイントをいくつかお伝えしたいと思います。4つありまして：まず1つ目は、HyperPodがインフラストラクチャ管理にかかる時間を削減できるということです。2つ目は、HyperPodは障害に対して回復力があり、自動修復機能を備えているため、手動で対処する必要がないということです。3つ目は、様々なツールに対する選択の柔軟性を提供するということ、そして4つ目は、コストを大幅に削減できるということです。
最後に、いくつかのリソースをご紹介させていただきます。このスライドの写真を撮っていただければと思います。私たちが提供している多くの機能に関するアナウンスメントがあり、また、セルフサービス形式でお試しいただけるワークショップもご用意しています。これらのワークショップでは、実際に導入を決める前に、様々なツールや選択肢を試していただくことができます。本日は皆様、お時間をいただき、ありがとうございました。残りのre:Inventも素晴らしいものになりますように。
※ こちらの記事は Amazon Bedrock を利用することで全て自動で作成しています。
※ 生成AI記事によるインターネット汚染の懸念を踏まえ、本記事ではセッション動画を情報量をほぼ変化させずに文字と画像に変換することで、できるだけオリジナルコンテンツそのものの価値を維持しつつ、多言語でのAccessibilityやGooglabilityを高められればと考えています。
Discussion