Amazon Web Services ブログ
大規模な分散トレーニングに特化したインフラストラクチャ、Amazon SageMaker HyperPod のご紹介
11月29日は、Amazon SageMaker HyperPod を紹介します。この製品は、大規模な分散トレーニング専用のインフラストラクチャを提供することで、基盤モデル (FM) のトレーニング時間を短縮するのに役立ちます。SageMaker がクラスターの状態をアクティブに監視し、障害のあるノードを交換してチェックポイントからモデルトレーニングを再開することで、ノードとジョブの回復を自動化しながら、SageMaker HyperPod を使用して FM を数週間から数か月間トレーニングできるようになりました。
クラスターには、SageMaker の分散トレーニングライブラリがあらかじめ設定されています。これにより、トレーニングデータとモデルをすべてのノードに分割して並列処理し、クラスタのコンピューティングとネットワークインフラストラクチャを最大限に活用できます。追加のフレームワーク、デバッグツール、最適化ライブラリをインストールすることで、トレーニング環境をさらにカスタマイズできます。
SageMaker HyperPod を使い始める方法を紹介します。次のデモでは、SageMaker HyperPod を作成し、AWS ML トレーニングリファレンスアーキテクチャ GitHub リポジトリで共有されている例を使用して Llama 2 7B モデルをトレーニングする方法を示します。
クラスターの作成と管理
SageMaker HyperPod 管理者は、AWS マネジメントコンソールまたは AWS コマンドラインインターフェイス (AWS CLI) を使用してクラスターを作成および管理できます。 コンソールで、Amazon SageMaker に移動し、左側のメニューの[HyperPod クラスター]で[クラスター管理]を選択し、その後[クラスターの作成]を選択します。
以下のセットアップでは、クラスター名を指定し、選択したインスタンスタイプと各インスタンスグループに割り当てるインスタンスの数でインスタンスグループを設定します。
また、クラスターの作成時に各インスタンスグループで実行するライフサイクルスクリプトを 1 つ以上準備して Amazon Simple Storage Service (Amazon S3) バケットにアップロードする必要があります。ライフサイクルスクリプトを使用すると、クラスター環境をカスタマイズし、必要なライブラリとパッケージをインストールできます。SageMaker HyperPod のライフサイクルスクリプトの例は、GitHub リポジトリにあります。
AWS CLI を使用する
AWS CLI を使用してクラスターを作成および管理することもできます。このデモでは、JSON ファイルでクラスター構成を指定します。2 つのインスタンスグループを作成することにしました。1 つは「controller-group」と呼ばれるクラスターコントローラーノード用で、もう 1 つは「worker-group」と呼ばれるクラスターワーカーノード用です。 モデルトレーニングを行うワーカーノードには、AWS Trainium チップを搭載した Amazon EC2 Trn1 インスタンスを指定します。
クラスターを作成するには、次の AWS CLI コマンドを実行します。
作成時に、aws sagemaker 記述クラスター
と aws sagemaker リストクラスターノード
を使用して、クラスターとノードの詳細を表示できます。コントローラーノードのクラスター ID とインスタンス ID を書き留めます。クラスターに接続するには、その情報が必要です。
また、Amazon FSx for Lustre などの共有ファイルシステムをアタッチすることもできます。FSx for Lustre を使用するには、Amazon Virtual Private Cloud (Amazon VPC) 設定を使用してクラスターをセットアップする必要があります。これは SageMaker VPC を作成する方法と Lustre に FSx をデプロイする方法を示す AWS CloudFormation テンプレートです。
クラスターに接続するには、
クラスターユーザーは、クラスター管理者によってプロビジョニングされたクラスターにアクセスできる必要があります。アクセス権限を設定したら、SSH を使用してクラスターに接続し、ジョブをスケジュールして実行できます。プリインストールされている AWS Systems Manager 用の AWS CLI プラグインを使用して、クラスターのコントローラノードに接続できます。
デモでは、コントロールノードのクラスター ID とインスタンス ID をターゲットとして指定して、次のコマンドを実行します。
Slurm を使用してクラスターでジョブをスケジュールして実行する
発売時には、SageMaker HyperPod は Slurm によるワークロードオーケストレーションをサポートしています。Slurm は人気のあるオープンソースのクラスター管理およびジョブスケジューリングシステムです。クラスター作成の一環として、ライフサイクルスクリプトを使用して Slurm をインストールして設定できます。ライフサイクルスクリプトの例はその方法を示しています。次に、標準の Slurm コマンドを使用してジョブをスケジュールし、起動できます。アーキテクチャの詳細と役立つコマンドについては、Slurm クイックスタートユーザーガイドをご覧ください。
このデモでは、AWS ML トレーニングリファレンスアーキテクチャ GitHub リポジトリにあるこの例を使用して、Trn1 インスタンスを使用して Slurm で Llama 2 7B をトレーニングする方法を示しています。私のクラスターはすでに Slurm でセットアップされており、Lustre ファイルシステムの FSx がマウントされています。
注意
Llama 2 モデルは Meta によって管理されています。Meta リクエストアクセスページからアクセスをリクエストできます。
クラスター環境のセットアップ
SageMaker HyperPod は、Conda、venv、Docker、enroot など、さまざまな環境でのトレーニングをサポートしています。README の指示に従って、仮想環境 aws_neuron_venv_pytorch
を構築し、Trn1 インスタンスでモデルをトレーニングするための torch_neuronx と neuronx-nemo-megatron ライブラリをセットアップしました。
モデル、トークナイザー、データセットの準備
指示に従って Llama 2 モデルとトークナイザーをダウンロードし、モデルを Hugging Face 形式に変換します。次に、RedPajama データセットをダウンロードしてトークン化します。最後の準備ステップとして、モデルトレーニングをスピードアップするために、事前 (AOT) コンパイルを使用して Llama 2 モデルをプリコンパイルします。
クラスターでジョブを起動する
これで、sbatch
コマンドを使用してモデルトレーニングジョブを開始する準備ができました。
squeue
コマンドを使用してジョブキューを表示できます。トレーニングジョブが実行されると、SageMaker HyperPod の復元機能が自動的に有効になります。SageMaker HyperPod は、前述のコマンドに示すように、ハードウェア障害を自動的に検出し、必要に応じてノードを交換し、自動再開
パラメーターが設定されている場合はチェックポイントからトレーニングを再開します。
モデルトレーニングジョブの出力は、次のファイルで確認できます。
モデルトレーニングが開始されたことを示すサンプル出力は、次のようになります。
モデルトレーニングジョブをさらに監視してプロファイリングするには、SageMaker がホストする TensorBoard やその他の任意のツールを使用できます。
今すぐご利用いただけます
SageMaker HyperPod は、本日より AWS リージョンの米国東部 (オハイオ州)、米国東部 (バージニア州北部)、米国西部 (オレゴン州)、アジアパシフィック地域 (シンガポール)、アジア太平洋地域 (シドニー)、アジアパシフィック地域 (東京)、欧州 (フランクフルト)、欧州 (アイルランド)、欧州 (ストックホルム) でご利用いただけます。
詳細はこちら。
- 価格情報とサポートされているクラスターインスタンスタイプのリストについては、Amazon SageMaker HyperPod を参照してください
- デベロッパーガイドをご覧ください
- AWS マネジメントコンソールにアクセスして、SageMaker HyperPod で FM のトレーニングを開始してください
– Antje
PS: AWS でブログ記事を書くのは、投稿タイトルの下に名前が 1 つしかない場合でも、常にチームの努力が必要です。今回は、Brad Doran、Justin Pirtle、Ben Snyder、Pierre-Yves Aquilanti、Keita Watanabe、Verdi March の各氏に、サンプルコードの提供や、大規模モデル学習インフラ、Slurm、SageMaker HyperPod の管理に関する専門知識の共有など、惜しみない協力をいただきました。
原文はこちらです。