Amazon Search M5 が AWS Trainium を使用して LLM トレーニングコストを 30% 節約した方法 | アマゾン ウェブ サービス

Amazon は何十年にもわたって機械学習 (ML) を開拓し、革新し、顧客に楽しい体験を提供してきました。 Amazon は初期の頃から、書籍の推奨、検索、不正行為の検出などのさまざまなユースケースに ML を使用してきました。 他の業界と同様に、加速されたハードウェアの進歩により、Amazon チームはニューラル ネットワークとディープ ラーニング (DL) を使用したモデル アーキテクチャを追求できるようになりました。

Amazon Search 内の M5 プログラムは、Amazon の発見学習戦略を所有し、多言語、マルチロケール、マルチエンティティ、マルチタスク、およびテキスト、画像、ビデオなどのマルチモーダルにわたる大規模なモデルを構築します。 M5 プログラムは、コストの最適化に対する厳格な管理を維持しながら、Amazon 全体の何百もの ML チームにユニバーサル エンベディングと大規模な基礎モデルを提供してきました。 これを達成するために、M5 チームはコストを削減するための新しい技術を定期的に評価しています。

多くの ML 組織と同様に、アクセラレータは主に DL トレーニングと推論を高速化するために使用されます。 AWS が最初のリリースで専用アクセラレータを立ち上げたとき AWSインフェレンティア 2020 年、M5 チームはすぐに それらを利用して、実稼働ワークロードをより効率的に展開します。、コストを節約し、待ち時間を短縮します。 昨年、AWS は AWS トレーニング アクセラレータは、次世代 DL モデルの開発および構築のコストあたりのパフォーマンスを最適化します。 この投稿では、M5 がモデルのトレーニングにかかる​​コストを 30% 削減できた方法について説明し、その過程で学んだベスト プラクティスのいくつかを共有します。

Trainium インスタンス

専用アクセラレータの進歩に伴い、Amazon は AWS Inferentia および Trainium の形で魅力的なアクセラレータも提供しています。 名前が示すように、これらのチップはそれぞれ推論とトレーニングのワークロードのニーズを超えるように最適化されています。 サイズが数十億パラメータに達する基礎モデルの大規模トレーニングには、Trainium Trn1 インスタンスと Trn1n インスタンス はその特性から理想的な選択です。 Trn1 インスタンスは最先端のテクノロジーを搭載しています ニューロンコア-v2、大量のアクセラレータ コンピューティングとメモリを備えています。 Trn1n インスタンスは、より大きなネットワーク帯域幅 (1,600 Gbs) 用に選択することもできるため、コストの最適化を念頭に置いたパフォーマンスの高いトレーニングに最適です。

アクセラレータを使用するには、アクセラレータをサポートするソフトウェア層が必要です。 Trn チップと Inf チッ​​プを使用すると、 AWS ニューロン SDK PyTorch XLA を利用して、Amazon 専用アクセラレータのロックを解除します。 PyTorch XLA は、PyTorch の Eager モードを Lazy モードのグラフベースの実装に変換します。 これらのグラフはその後使用され、アクセラレータで使用できるようにさらにコンパイルされます。 PyTorch Neuron (Neuron SDK の一部) を使用すると、PyTorch ユーザーは数行のコードで Trainium NeuronCore でモデルをトレーニングできます。

モデルとワークロード

M5 チームは、Amazon 全体のさまざまなチームがユーザーに喜びをもたらすのを支援するために、基礎モデルと普遍的な表現をトレーニングおよびデプロイします。 Amazon.com 顧客。 そのようなモデルの 1 つは、テキスト エンコーダー モデルと、その後に数億のトレーニング可能なパラメーターを備えたニューラル ネットワーク アーキテクチャによって定義された明示的または暗黙的な特徴相互作用を備えた多層パーセプトロン (MLP) です。 このモデルは数十億のトークンでトレーニングされ、オフライン バッチ推論設定で数百万の埋め込みを生成するために使用されます。 これらの埋め込みは、顧客向けの Tier-XNUMX Amazon サービスへの入力です。

実稼働パイプラインのインフラストラクチャでは、 AWSバッチ 　 フェアシェアキューイング戦略、EFA 対応のマルチノード trn1.32xlarge クラスターをモデル トレーニングのコンピューティングとして使用します。 機能的には、実稼働パイプラインは、基礎となる DL ライブラリとして PyTorch を使用して、増分モデル トレーニング、トレーニング済みモデルの評価、トレーニング済みモデルのオフライン バッチ推論を実行します。

目標

お客様に喜んでいただくことが第一の信条です。 パイプラインの顧客対応の性質を考慮すると、すべてのサービス レベル アグリーメント (SLA) が後退することなく満たされることが重要です。 私たちは、既存の GPU プロダクション パイプラインを適応させて Trainium に移行するための XNUMX つの重要な受け入れ基準を特定しました。

• モデルの品質 – モデルの品質は顧客エクスペリエンスに直接影響します。 GPU と Trainium の間のモデル品質の差が 0.1% 未満であることが必要です。
• トレーニングのスループット – お客様に最も新鮮なエクスペリエンスを提供するために、モデルを定期的に繰り返しトレーニングします。 実稼働 SLA を満たすためには、事前定義された期間 (1 週間など) 内にモデルの収束を達成する必要があります。

次のセクションでは、この基準から逆算して取り組んだ過程と、Amazon 規模の本番ワークロードをサポートするための学習について共有します。

トレーニングスクリプト

モデルのトレーニングを開始する前に、トレーニング スクリプトを XLA 準拠にするように変更する必要があります。 モデルのサイズを考慮して、分散データ並列 (DDP) を使用してモデルをトレーニングします。 DDP を使用すると、コードを変更することなく、モデル トレーニングの実行に使用するマシンの数をスケールアップすることで、モデル トレーニングのスループットを向上させることができます。 に記載されている指示に従いました。 Neuron PyTorch MLP トレーニング チュートリアル XLA 固有の構成要素をトレーニング スクリプトに追加します。 これらのコード変更は実装が簡単です。 以下は、モデルのスループットを大幅に向上させる演習から得られた重要な技術的な学習です。

• xm.mark_step() の配置 – xm.mark_step() 遅延収集された計算グラフをコンパイルして実行します。 呼び出し中 mark_step 呼び出し回数が多すぎると、小さなグラフが多数生成されますが、呼び出し回数が少なすぎると、生成されるグラフの数は少なくなりますが、大きなグラフが生成されます。 アプリケーションに応じて、モデル トレーニングのスループットと実装は、 xm.mark_step()。 私たちの実装では、 xm.mark_step() 前方パスと後方パスの後、およびオプティマイザー ステップの後。
• XLA マルチプロセッシング デバイス ローダーを使用したデータ ローダー ラッピング – これは見逃しがちな重要なステップです。 マルチプロセッシング デバイス ローダー torch_xla.distributed.parallel_loader.MpDeviceLoader スループットを向上させるために、データのロードをプリロードし、デバイスの実行とオーバーラップするオプションを使用して、各 XLA デバイスにトレーニング データをロードします。 デバイスローダーはまた、 xm.mark_step() したがって、ホストからデバイスにデータをロードするためのグラフを構築できます。

Trainium のコンピレーション

従来、GPU を使用したモデル開発サイクルには、モデルまたはトレーニング スクリプトに変更を加え、それを GPU デバイス上で直接実行することが含まれていました。 XLA を使用する Trainium などのアクセラレータでは、アクセラレータでモデル トレーニングを実行する前に追加の手順が必要です。 XLA 計算グラフは、コンパイルされた後にのみ実行できます。 一般に、このコンパイルを実行するには XNUMX つの方法があります。XNUMX つは最初にすべてのグラフをトレースおよびコンパイルしてから実行する Ahead of Time (AOT)、もう XNUMX つはグラフをトレースおよびコンパイルして実行する Just In Time (JIT) です。遭遇します。 Neuron SDK は、これらの両方をすぐに提供します。 通常、AOT コンパイルが最初に実行されます。 このコンパイル後にグラフが実行されます。 新しいグラフが見つかった場合、Neuron ランタイムはそれらを実行する前に JIT コンパイルを呼び出します。 AOT コンパイルを実行するために、Neuron SDK は以下を提供します。 ニューロン_パラレル_コンパイル、トレーニング スクリプトの試行実行からグラフを抽出し、並列 AOT コンパイルを実行するコンパイル ユーティリティです。

AOT コンパイルの重要な側面は、トレーニング中に新しい計算グラフが作成されないようにすることです。 新しい計算グラフ (したがって再コンパイル) のソースの 4 つは、モデルのトレーニング中のトレーニング バッチの動的形状です。 静的シェイプと固定サイズのバッチを使用すると、トレーニング時間のコンパイルが不要になり、モデルの精度に影響を与えることなくトレーニングのスループットが大幅に向上することがわかりました。 トレーニングにこのような制約を適用することで、AOT コンパイル中にすべてのグラフをトレースするには、モデル トレーニングの 5 ～ XNUMX ステップ、モデル検証の XNUMX ステップ、およびモデルのチェックポイント設定を XNUMX 回行うだけで済むことがわかりました。 Neuron SDK は常に進化しており、将来的には動的シェイプもサポートする予定であることに注意してください。

さらに、編集されたグラフは次の場所に保存されます。 ニューロン永続キャッシュ ディスク上または Amazon シンプル ストレージ サービス (Amazon S3) バケット。 これは、モデル アーキテクチャとトレーニング構成が変更されない運用ワークロードで特に役立ちます。 したがって、コンパイルのオーバーヘッドは XNUMX 回だけ発生します。 キャッシュの使用は、環境フラグを設定するのと同じくらい簡単です。

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Neuron コンパイラには、次の XNUMX つも提供されます。 コンパイラレベルの最適化オプション (O1、O2、O3) は、コンパイル時間とモデル実行スループットのバランスをとります。 O1 は計算グラフのコア最適化を有効にしてコンパイル時間を最小限に抑えます。O3 はコンパイル時間の増加と引き換えにモデル実行のスループットを向上させます。O2 (デフォルト オプション) は 1 つの間のバランスです。 このユースケースでは、O86 最適化を使用し、モデル精度メトリクスに変更を加えることなくコンパイル時間が 5% 短縮された一方、デフォルトの最適化 (O7) と比較してスループットが約 2 ～ XNUMX% 削減されたことがわかりました。 ユースケースに応じて、さまざまなレベルの最適化を選択できます。

要約すると、コンパイルには次のフラグを使用しました。

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

チェックポイントの互換性

コンパイルが正常に完了したら、Trainium でのモデルのトレーニングに進むことができます。 前述したように、モデルは段階的にトレーニングされます。つまり、以前にトレーニングされたモデルのチェックポイントをロードし、新しいデータでトレーニングを継続します。 PyTorch と PyTorch XLA を使用すると、チェックポイントの相互運用性を通じてアクセラレータ間のシームレスな移行が可能になります。 GPU と Trainium の間を柔軟に移動できるため、以前の GPU モデルをシームレスにロードして Trainium マシンでトレーニングできるようになりました。 これは、本番環境のダウンタイムやモデルの精度の低下なしに、事前にトレーニングされた最良のモデルでモデルを初期化できるようにするために重要でした。

GPU モデルは標準の PyTorch モデル保存ユーティリティを使用して保存されたため、PyTorch チェックポイント読み込みユーティリティを使用して GPU モデルを Trainium デバイスにロードすることができました。

たとえば、GPU/CPU では、次のコードを使用してモデルを保存できます。

torch.save(model.state_dict(), PATH)

次に、モデルを Trainium にロードし直します。

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

同様に、次のコードを使用してモデルを Trainium に保存できます。

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

そして、モデルを GPU/CPU にロードし直します。

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

実際、モデルのトレーニングに DDP を使用しているため、モデルの読み込みは前のチェックポイントのトレーニングに使用されたマシンの数に依存しません。 これにより、コードの変更やモデルのトレーニングへの悪影響を与えることなく、Trn1 フリートを水平方向にスケーリングすることができます。 これらの PyTorch ベースのチェックポイントは、AWS Inferentia2 またはその他のアクセラレータでの推論ユースケースに直接使用することも、トーチスクリプト化することもできます。

動作の安定性

実稼働環境でワークロードを実行するには、複数の SLA を満たす必要があることは、どれだけ強調してもしすぎることはありません。 私たちのユースケースでは、モデルの品質とトレーニング スループットの SLA とは別に、運用パイプラインが運用的に安定していることが不可欠であり、これはモデルのトレーニング、評価、推論中のダウンタイムや中断を最小限に抑えることを意味します。

既存の GPU ベースのパイプラインと同様に、パイプラインの動作を安定させるために多数のメカニズムを追加しました。 モデルのトレーニングを開始する前に、複数の健全性テストを実行してマシンの健全性を評価します。 これらのテストには通常、アクセラレータ デバイスの健全性を検証するための単純なテンソル操作が含まれます。 分散トレーニングでは、インスタンス間の集団通信を検証するテストを実行することも重要であることがわかりました。 私たちが使用したのは、 NCCOM テストスイート Neuron SDK からこれを実現するために、all-gather、all-reduce、reduce-scatter などのさまざまな操作を実行します。

これまで述べた提案に従った後でも、基盤となるアクセラレータに関係なく、どのパイプラインでも一時的な問題が避けられないことがわかりました。 トレーニング パイプラインの復元力を構築するには、これらの潜在的な問題を解決するための再試行メカニズムを組み込むことをお勧めします。 を使用しております AWS Batch の自動再試行 モデルのトレーニング中に一時的なエラーが発生したジョブを再試行します。 トレーニングの終わり頃に障害が発生した場合、これらの再起動にはコストがかかる可能性があります。 この問題に対処するために、以前にトレーニングされたモデルのチェックポイントをロードし、その時点からトレーニングを継続するようにトレーニング スクリプトを調整しました。 この機能を使用すると、失敗したトレーニング ジョブを最小限のオーバーヘッドで積極的に再開できます。

これらの復元メカニズムを導入することで、Trn98.5 上のワークロードの成功率 1% を達成することができました。これは、既存の GPU パイプラインの成功率に匹敵します。

結果

モデルの精度を検証するために、同じ GPU チェックポイントから 0.1 つのモデルを初期化し、XNUMX つを Trainium でトレーニングし、もう XNUMX つを同等の GPU でトレーニングしました。 どちらのモデルも、同じトレーニング ハイパーパラメータを使用してトレーニングされました。 メトリクスの計算に使用されるデータセットはホールドアウト データセットであり、このデータセットに対するモデルの精度を N グローバル ステップごとに評価します。 X 軸はグローバル ステップ、Y 軸はモデルの精度です。 次のグラフの各点でモデルの精度の差が XNUMX% 未満であることが観察されました。

さらに、モデル トレーニングの費用対効果を評価するために、モデルの収束に達するまでにかかった実時間を比較することを好みます。 これにより、トークンあたりのコスト、達成された FLOPS/ドル、その他の要素などの指標と比較して、コスト削減のより実用的なビューが得られると考えています。 trn1.32xlと同等のトレーニング時間を考慮すると アマゾン エラスティック コンピューティング クラウド (Amazon EC2) インスタンスの場合、Trainium ではモデルの収束にかかるコストが最大 30% 低くなることがわかりました。

まとめ

DL ワークロードに応じてさまざまなアクセラレータを評価する際には、考慮すべき要素が数多くあります。 最も重要なものには、モデルの品質、スループット、コスト、可用性があります。 選択したアクセラレータに基づいてモデルの品質とスループットが犠牲にならないようにすることが最も重要です。

Annapurna Neuron チームとのパートナーシップと協力のおかげで、Amazon Search M5 チームは Trainium に移行することでコストを最大 30% 節約することができました。 チームは Trainium を使用して、市場の同等のアクセラレータと同等のモデル品質とスループットを達成できます。 XLA のサポートによるチェックポイントの相互運用性と最小限のコード変更により、M5 はワークロードに応じて複数のアクセラレータから選択できるようになりました。 これにより、M5 チームは Trainium の大規模な計算能力を活用し、Amazon.com の顧客を満足させるアクセラレータに依存しないソリューションを構築できるようになりました。 運用の観点から見ると、Trainium は Amazon 規模の Tier 1 サービスをサポートできることが証明されています。 M5 チームは、Amazon に最低コストで最高のモデルを提供するために、より多くのワークロードを Trainium に移行し続けています。

要約すると、M5 チームはアクセラレーター群に Trainium を追加することで、コスト効率の高い実稼働グレードの ML トレーニングを実行できるようになりました。 ML ワークロード専用の Amazon シリコンのメリットを享受するには、Trainium や AWS Inferentia などの他の Neuron デバイスを検討することをお勧めします。 さまざまなモデルを紹介する多数のチュートリアルの XNUMX つから簡単に始めましょう。 ラマ 2、Trainium で利用可能.

著者について

アビナンダン・パトニ Amazon Search のシニア ソフトウェア エンジニアです。 彼は、スケーラブルな分散型深層学習トレーニングとリアルタイム推論のためのシステムとツールの構築に焦点を当てています。

ジェームズ・パーク アマゾン ウェブ サービスのソリューション アーキテクトです。 彼は Amazon.com と協力して AWS 上のテクノロジー ソリューションを設計、構築、デプロイしており、特に AI と機械学習に興味を持っています。 余暇には、新しい文化、新しい経験を探し、最新のテクノロジー トレンドを把握することを楽しんでいます。 あなたは彼を見つけることができます LinkedIn.

ジェリー・マンニル Amazon Search のソフトウェア エンジニアです。 彼は、分散トレーニング インフラストラクチャの効率、堅牢性、拡張性の向上に取り組んでいます。

ケン・スー Amazon Search のソフトウェア エンジニアです。 彼は、トレーニングの効率とスケーラブルな分散トレーニング ワークフローの改善に取り組んでいます。 仕事以外では、ハイキングとテニスが好きです。

RJ Amazon社内のエンジニアです。 彼は、トレーニング用の分散システムのシステムを構築および最適化し、ML 推論のレイテンシーを削減するためにシステムの導入を最適化することに取り組んでいます。 仕事以外では、食品レシピの作成に生成 AI を使用することを検討しています。

• SEO を活用したコンテンツと PR 配信。 今日増幅されます。
• PlatoData.Network 垂直生成 Ai。 自分自身に力を与えましょう。 こちらからアクセスしてください。
• プラトアイストリーム。 Web3 インテリジェンス。 知識増幅。 こちらからアクセスしてください。
• プラトンESG。 カーボン、 クリーンテック、 エネルギー、 環境、 太陽、 廃棄物管理。 こちらからアクセスしてください。
• プラトンヘルス。 バイオテクノロジーと臨床試験のインテリジェンス。 こちらからアクセスしてください。
• 情報源： https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/