『Amazon Search M5がAWS Trainiumを使用してLLMトレーニングコストを30%節約しました』

「Amazon Search M5がAWS Trainiumを活用し、LLMトレーニングコストを30%節約！」

数十年にわたり、Amazonは機械学習（ML）の先駆者であり、革新をもたらし、顧客に楽しい体験を提供してきました。最初の日々から、Amazonは書籍の推薦、検索、詐欺検出など、さまざまなユースケースでMLを使用してきました。業界全体と同様に、高速ハードウェアの進歩により、Amazonチームはニューラルネットワークと深層学習（DL）を使用したモデルアーキテクチャの追求が可能になりました。

Amazon Search内のM5プログラムは、Amazonの検索における探索的学習戦略を担当し、テキスト、画像、ビデオなどのマルチモーダル（多言語、多地域、多エンティティ、マルチタスクなど）の大規模モデルを構築しています。M5プログラムは、数百のMLチームに対してユニバーサルな埋め込みと大規模な基盤モデルを提供しており、コスト最適化に厳格な制御を保ちながらサービスを提供しています。そのため、M5チームでは定期的に新しいテクニックを評価してコストを削減しています。

多くのML組織と同様に、アクセラレータは主にDLのトレーニングと推論を加速するために使用されています。AWSが2020年に最初のリリースでAWS Inferentiaを導入した際、M5チームはすぐにそれらを効率的に展開して生産ワークロードを実行するために利用し始めました。これにより、コストを削減し、レイテンシーを減らすことができました。昨年、AWSはAWS Trainiumアクセラレータを導入し、次世代のDLモデルの開発と構築においてコストパフォーマンスを最適化しています。この記事では、M5がモデルのトレーニングコストを30%削減することに成功した方法と、その過程で学んだいくつかのベストプラクティスについて説明します。

Trainiumインスタンス

目的に特化したアクセラレータの進化とともに、AmazonはAWS InferentiaとTrainiumという魅力的なアクセラレータも提供しています。それぞれの名前が示すように、これらのチップは推論とトレーニングのワークロードのニーズを超えるように最適化されています。数十億のパラメータを持つ基盤モデルを大規模にトレーニングする場合、TrainiumのTrn1およびTrn1nインスタンスが特に適しています。Trn1インスタンスは最新のNeuronCore-v2を搭載しており、アクセラレータの計算およびメモリが豊富にあります。Trn1nインスタンスはネットワーキング帯域幅（1,600 Gb）がより大きいため、コスト最適化を考慮した高性能なトレーニングに適しています。

• カスタムレンズを使用して、優れたアーキテクチャのIDPソリューションを構築する – パート5：コスト最適化
• 自分自身のレンズでウェルアーキテクチャなIDPソリューションを構築する – パート6：持続可能性
• 「カスタムレンズを使用して、よく設計されたIDPソリューションを構築する-パート4パフォーマンス効率性」

アクセラレータを使用するためには、それをサポートするソフトウェアレイヤーが必要です。TrnとInfチップでは、AWS Neuron SDKを使用して、PyTorch XLAのサポートによりAmazonの目的に特化したアクセラレータを利用することができます。PyTorch XLAはPyTorchのイーガーモードをレイジーモードのグラフベースの実装に変換します。これらのグラフはアクセラレータと一緒に使用され、さらにコンパイルされます。PyTorch Neuron（Neuron SDKの一部）を使用することで、PyTorchユーザーはわずかなコードでTrainium NeuronCoresでモデルをトレーニングすることができます。

モデルとワークロード

M5チームは、Amazon全体の様々なチームがAmazon.comの顧客に喜びをもたらすために、基礎モデルとユニバーサル表現をトレーニングおよび展開しています。そのようなモデルの1つは、テキストエンコーダーモデルであり、ニューラルネットワークアーキテクチャによって定義された数億以上のトレーニング可能なパラメータを持つ多層パーセプトロン（MLP）です。このモデルは数十億のトークンでトレーニングされ、オフラインのバッチ推論設定で数百万の埋め込みを生成するために使用されます。これらの埋め込みは顧客のフロントエンドのTier-1 Amazonサービスへの入力です。

プロダクションパイプラインのインフラストラクチャでは、AWS Batchとフェアシェアキューイングストラテジーを使用し、モデルトレーニングの計算にEFA対応のマルチノードtrn1.32xlargeクラスターを使用しています。機能的には、プロダクションパイプラインはインクリメンタルなモデルトレーニング、トレーニングされたモデルの評価、およびトレーニングされたモデルのオフラインバッチ推論をPyTorchを基に行っています。

ゴール

お客様を喜ばせることが最も重要な原則です。パイプラインの顧客対応性を考慮すると、すべてのサービスレベル契約（SLA）を回帰なしで満たすことが重要です。私たちは、既存のGPU生産パイプラインを適応させてTrainiumに移行するための2つの重要な受け入れ基準を特定しました：

• モデルの品質 – モデルの品質は直接顧客の体験に影響します。GPUとTrainiumの間のモデル品質の違いは0.1％未満である必要があります。
• トレーニングのスループット – 定期的にモデルを反復的にトレーニングして、顧客に最新の体験を提供します。モデルの収束は、あらかじめ定義された期間（たとえば1週間）内に達成されなければ、製品のSLAを満たせません。

以下のセクションでは、この基準から逆算して作業を進めた結果と、Amazonスケールのプロダクションワークロードをサポートするための学びを共有します。

トレーニングスクリプト

モデルのトレーニングを開始する前に、トレーニングスクリプトに変更を加えてXLAに準拠させる必要があります。モデルのサイズに応じて、分散データ並列（DDP）を使用してモデルをトレーニングしています。DDPを使用することで、モデルのトレーニングスループットを向上させるために、モデルのトレーニングを実行するために使用するマシンの数をスケーリングアップすることができます。コードの変更は必要ありません。私たちは、Neuron PyTorch MLPトレーニングチュートリアルで提供されている指示に従って、トレーニングスクリプトにXLA固有の構造を追加しました。これらのコードの変更は、簡単に実装できます。以下は、私たちのモデルスループットを大幅に向上させたこの練習から得たいくつかの重要な技術的な学びです：

• xm.mark_step()の配置 – xm.mark_step()は、遅延評価された計算グラフをコンパイルして実行します。mark_stepを頻繁に呼び出すと、小さなグラフの数が増えますが、呼び出し回数が少ないと、大きなグラフの数が増えます。アプリケーションによっては、xm.mark_step()の配置によってモデルのトレーニングのスループットや実装が異なることがあります。私たちの実装では、forwardおよびbackwardパスの後に1つのxm.mark_step()を配置し、optimizerステップの後にも1つ配置しています。
• XLA多重プロセスデバイスローダーでのデータローダーのラッピング – これは見落とされる可能性のある重要なステップです。マルチプロセスデバイスローダーtorch_xla.distributed.parallel_loader.MpDeviceLoaderは、トレーニングデータをXLAデバイスにロードします。スループットを向上させるためのデバイス実行との重なりを提供することもできます。デバイスローダーはまた、xm.mark_step()を呼び出すので、ホストからデバイスへのデータロードのためのグラフを構築することができます。

Trainiumのためのコンパイル

従来のGPUを使用したモデル開発サイクルでは、モデルまたはトレーニングスクリプトに変更を加え、直接GPUデバイスで実行します。XLAを使用するTrainiumなどのアクセラレータは、モデルトレーニングを実行する前に追加のステップが必要です。XLAの計算グラフはコンパイルされた後に実行できます。一般的に、このコンパイルを実行するためには2つの方法があります：Ahead of Time（AOT）では、まずすべてのグラフを追跡してコンパイルし、それから実行します。Just In Time（JIT）では、グラフはトレースされ、コンパイルされ、その場で実行されます。Neuron SDKはこれらの両方を提供しています。通常、まずAOTコンパイルが実行されます。その後、グラフが実行されます。新しいグラフが出現する場合、Neuronランタイムはそれらを実行する前にJITコンパイルを実行します。AOTコンパイルを実行するために、Neuron SDKはneuron_parallel_compileというコンパイルユーティリティを提供しています。このユーティリティは、トレーニングスクリプトの試行実行からグラフを抽出し、並列AOTコンパイルを実行します。

AOTコンパイルの重要な側面は、トレーニングの過程で新しい計算グラフが作成されないようにすることです。モデルトレーニング中のトレーニングバッチのダイナミックな形状は、新しい計算グラフ（および再コンパイル）の原因となります。静的な形状と固定サイズのバッチを使用すると、トレーニング時間のコンパイルがなくなり、モデルの精度に影響を与えることなくトレーニングスループットが大幅に向上することがわかりました。このようなトレーニングの制約を強制することで、AOTコンパイル中にすべてのグラフをトレースするためにモデルトレーニングのために4〜5ステップ、モデル検証のために1ステップ、モデルを1回チェックポイントする必要があることがわかりました。重要なことは、Neuron SDKが常に進化しており、将来的には動的な形状もサポートするということです。

さらに、コンパイルされたグラフはディスク内のニューロン持続キャッシュまたはAmazon Simple Storage Service（Amazon S3）バケットに保存されます。これは、モデルのアーキテクチャとトレーニング構成が変わらないプロダクションワークロードに特に便利です。そのため、コンパイルのオーバーヘッドは一度だけ発生します。キャッシュの使用は、環境フラグを設定するだけで簡単です：

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

ニューロンコンパイラは、コンパイル時間とモデルの実行スループットのバランスを取るための3つのコンパイラレベルの最適化オプション（O1、O2、O3）を提供しています。O1は計算グラフのコア最適化を有効にし、コンパイル時間を最小限に抑えます。O3はコンパイル時間を増やす代わりに、改良されたモデルの実行スループットを提供します。O2（デフォルトオプション）はその両方のバランスです。私たちのユースケースでは、O1の最適化を使用し、モデルの精度メトリックに変化はないまま、コンパイル時間が86％減少し、デフォルトの最適化（O2）と比較してスループットが約5-7％減少することを確認しました。ユースケースに応じて、さまざまなレベルの最適化を選択できます。

まとめると、次のフラグをコンパイルに使用しました：

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のチェックポイントローディングユーティリティを使用してTrainiumデバイスにGPUモデルをロードすることができました。

たとえば、GPU/CPUで次のコードでモデルを保存できます：

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

次に、Trainiumでモデルをロードします：

import torch_xla.core.xla_model as xmxla_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# 自動的にデータをCPUに移動し、マスターデバイス用にxm.save(model.state_dict(), PATH) 

そして、GPU/CPUで次のコードでモデルをロードします：

model = MyModel(*args, **kwargs)model.load_state_dict(torch.load(PATH))model.to(device) # どのデバイスでも構いません

実際には、モデルのトレーニングにはDDPを使用しているため、前のチェックポイントで使用されたマシンの数に関係なく、モデルのローディングは無関係です。これにより、Trn1フリートを水平にスケーリングすることができ、モデルのトレーニングにはコードの変更や副作用がありません。これらのPyTorchベースのチェックポイントは、AWS Inferentia2や他のアクセラレータ上での推論ユースケースに直接使用することもできます。

運用の安定性

プロダクションでのワークロードの実行には、複数のSLAの達成が必要です。私たちのユースケースでは、モデルの品質とトレーニングのスループットを除いても、モデルのトレーニング、評価、推論の際には最小限のダウンタイムと中断がありません。既存のGPUベースのパイプラインと同様に、パイプラインの運用安定性を確保するために、さまざまなメカニズムを追加しました。モデルのトレーニングを開始する前に、マシンの状態を評価するために複数のセーニティテストを実行します。これらのテストには、アクセラレータデバイスの状態を確認するためのシンプルなテンソル演算が含まれています。分散トレーニングでは、インスタンス間の集合通信を検証するためのテストも重要です。私たちは、Neuron SDKのNCCOMテストスイートを使用して、全体集合、全体削減、およびリデューススキャッターなどのさまざまな操作を実行してこれを達成しました。

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

これらの堅牢性メカニズムが備わっていることで、Trn1でのワークロードの成功率は私たちの既存のGPUパイプラインの成功率と同等の98.5%を達成することができました。

結果

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

さらに、モデルトレーニングのコスト効率を評価するために、モデル収束までにかかる壁時計時間を比較することを好みます。これにより、コストパーハイパラメータ、達成FLOPS/ドルなどの指標と比較するよりも、より実践的なコスト削減の視点を提供すると考えています。trn1.32xlと比較可能なAmazon Elastic Compute Cloud (Amazon EC2) インスタンスのトレーニング時間を考慮すると、Trainiumはモデル収束までのコストが最大30%削減できることが観察されました。

結論

DLワークロードの異なるアクセラレータを評価する際には、モデルの品質、スループット、コスト、利用可能性など、多くの要素を考慮する必要があります。選んだアクセラレータに基づいてモデルの品質とスループットが損なわれないようにすることが非常に重要です。

Annapurna Neuron チームとのパートナーシップと協力により、Amazon Search M5 チームはTrainiumへの移行により最大30%のコスト削減を実現することができました。チームはTrainiumを使用し、市場での比較可能なアクセラレータと同等のモデル品質とスループットを達成することができました。チェックポイントの互換性とXLAのサポートにより、M5はワークロードに適した複数のアクセラレータから選択することができました。これにより、M5チームはTrainiumの大規模な計算能力を活用し、Amazon.comの顧客に喜ばれるアクセラレータに依存しないソリューションを構築することができました。運用上の観点から、TrainiumはAmazonスケールでの第1クラスのサービスをサポートする能力が証明されています。M5チームは引き続き、最低のコストでAmazonのために最高のモデルを提供するために、さらに多くのワークロードをTrainiumに移行しています。

まとめると、M5チームはTrainiumをアクセラレータのフリートに追加することで、コスト効果の高い本格的なMLトレーニングを実施することができました。TrainiumやAWS InferentiaなどのNeuronデバイスを活用することで、目的に特化したAmazonのシリコンをMLワークロードに使用する利点を享受することをお勧めします。Trainiumで提供されるLlama 2などのさまざまなモデルを使用したチュートリアルを利用して簡単に始めることができます。

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