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 クラスター]で[クラスター管理]を選択し、その後[クラスターの作成]を選択します。

Amazon SageMaker HyperPod クラスター

以下のセットアップでは、クラスター名を指定し、選択したインスタンスタイプと各インスタンスグループに割り当てるインスタンスの数でインスタンスグループを設定します。

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 インスタンスを指定します。

// demo-cluster.json
{
   "InstanceGroups":[
        {
            "InstanceGroupName": "controller-group",
            "InstanceType": "ml.m5.xlarge",
            "InstanceCount": 1,
            "lifecycleConfig": {
                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
            "ThreadsPerCore": 1
        },
        {
            "InstanceGroupName": "worker-group",
            "InstanceType": "trn1.32xlarge",
            "InstanceCount": 4,
            "lifecycleConfig": {
                "SourceS3Uri": "s3://<your-s3-bucket>/<lifecycle-script-directory>/",
                "OnCreate": "on_create.sh"
            },
            "ExecutionRole": "arn:aws:iam::111122223333:role/my-role-for-cluster",
            "ThreadsPerCore": 1
        }
    ]
}

クラスターを作成するには、次の AWS CLI コマンドを実行します。

aws sagemaker create-cluster \
--cluster-name antje-demo-cluster \
--instance-groups file://demo-cluster.json

作成時に、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 をターゲットとして指定して、次のコマンドを実行します。

aws ssm start-session \
--target sagemaker-cluster:ntg44z9os8pn_i-05a854e0d4358b59c \
--region us-west-2

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 コマンドを使用してモデルトレーニングジョブを開始する準備ができました。

sbatch --nodes 4 --auto-resume=1 run.slurm ./llama_7b.sh

squeue コマンドを使用してジョブキューを表示できます。トレーニングジョブが実行されると、SageMaker HyperPod の復元機能が自動的に有効になります。SageMaker HyperPod は、前述のコマンドに示すように、ハードウェア障害を自動的に検出し、必要に応じてノードを交換し、自動再開パラメーターが設定されている場合はチェックポイントからトレーニングを再開します。

モデルトレーニングジョブの出力は、次のファイルで確認できます。

tail -f slurm-run.slurm-<JOB_ID>.out

モデルトレーニングが開始されたことを示すサンプル出力は、次のようになります。

Epoch 0:  22%|██▏       | 4499/20101 [22:26:14<77:48:37, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.43, v_num=5563, reduced_train_loss=2.470, gradient_norm=0.121, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.40]
Epoch 0:  22%|██▏       | 4500/20101 [22:26:32<77:48:18, 17.95s/it, loss=2.44, v_num=5563, reduced_train_loss=2.450, gradient_norm=0.120, parameter_norm=1864.0, global_step=4512.0, consumed_samples=1.16e+6, iteration_time=16.50]

モデルトレーニングジョブをさらに監視してプロファイリングするには、SageMaker がホストする TensorBoard やその他の任意のツールを使用できます。

今すぐご利用いただけます
SageMaker HyperPod は、本日より AWS リージョンの米国東部 (オハイオ州)、米国東部 (バージニア州北部)、米国西部 (オレゴン州)、アジアパシフィック地域 (シンガポール)、アジア太平洋地域 (シドニー)、アジアパシフィック地域 (東京)、欧州 (フランクフルト)、欧州 (アイルランド)、欧州 (ストックホルム) でご利用いただけます。

詳細はこちら。

– Antje

PS: AWS でブログ記事を書くのは、投稿タイトルの下に名前が 1 つしかない場合でも、常にチームの努力が必要です。今回は、Brad Doran、Justin Pirtle、Ben Snyder、Pierre-Yves Aquilanti、Keita Watanabe、Verdi March の各氏に、サンプルコードの提供や、大規模モデル学習インフラ、Slurm、SageMaker HyperPod の管理に関する専門知識の共有など、惜しみない協力をいただきました。

原文はこちらです。