本日、Amazon SageMaker HyperPod をご紹介します。これは、大規模な分散トレーニング専用のインフラストラクチャを提供することで、ベースモデルのトレーニング時間を短縮します。SageMaker HyperPod を使用すれば、数週間、あるいは数か月にわたってベースモデルをトレーニングできます。SageMaker はクラスターの健全性をプロアクティブに監視し、障害が発生したノードを置き換えてチェックポイントからモデルトレーニングを再開することで、ノードとジョブの復元力を自動的に提供します。

クラスターにはSageMakerの分散トレーニングライブラリがあらかじめ設定されており、トレーニングデータとモデルをすべてのノードに分割して並列処理することで、クラスターのコンピューティングおよびネットワークインフラストラクチャを最大限に活用できます。また、追加のフレームワーク、デバッグツール、最適化ライブラリをインストールして、トレーニング環境をさらにカスタマイズすることもできます。

以下は、SageMaker HyperPod を使い始めるためのガイドです。以下のデモでは、Amazon Machine Learning トレーニングリファレンスアーキテクチャ GitHub リポジトリで共有されているサンプルを使用して、Llama 2 7B モデルをトレーニングする方法を示すために、SageMaker HyperPod を作成しました。

クラスターの作成と管理

SageMaker HyperPod 管理者は、Amazon マネジメントコンソールまたは Amazon コマンドラインインターフェース (Amazon CLI) を使用してクラスターを作成および管理できます。コンソールでAmazon SageMakerに移動し、左側のメニューの「HyperPod Clusters」の下の「Cluster Management」を選択し、 「Create Cluster」を選択します。

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

さらに、1つ以上のライフサイクルスクリプトを準備し、Amazon Simple Storage Service (Amazon S3) バケットにアップロードして、クラスター作成時に各インスタンスグループで実行できるようにする必要があります。ライフサイクルスクリプトを使用すると、クラスター環境をカスタマイズし、必要なライブラリとパッケージをインストールできます。SageMaker HyperPodのサンプルライフサイクルスクリプトは、GitHubリポジトリで入手できます。

Amazon CLIの使用

Amazon CLI を使用してクラスターを作成および管理することもできます。デモでは、JSON ファイルでクラスター構成を指定しました。クラスターコントローラーノード用の「controller-group」と、クラスターワーカーノード用の「worker-group」という 2 つのインスタンスグループを作成することにしました。モデルのトレーニングを実行するワーカーノードには、Amazon Trainium チップを搭載した Amazon EC2 Trn1 インスタンスを指定しました。

 // デモクラスター.json[
{
"InstanceGroupName": "コントローラグループ",
「インスタンスタイプ」: 「ml.m5.xlarge」
「インスタンス数」: 1,
「ライフサイクル設定」: {
"SourceS3Uri": "s3://<S3バケット>/<ライフサイクルスクリプトディレクトリ>/",
"OnCreate": "on_create.sh"
},
"実行ロール": "arn:aws:iam::111122223333:role/my-role-for-cluster",
「コアあたりのスレッド数」: 1
},
{
"InstanceGroupName": "ワーカーグループ",
「インスタンスタイプ」: 「ml.trn1.32xlarge」
「インスタンス数」: 4,
「ライフサイクル設定」: {
"SourceS3Uri": "s3://<S3バケット>/<ライフサイクルスクリプトディレクトリ>/",
"OnCreate": "on_create.sh"
},
"実行ロール": "arn:aws:iam::111122223333:role/my-role-for-cluster",
「コアあたりのスレッド数」: 1
}] 

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

 aws sagemaker クラスター作成 \
    --クラスター名 アンチデモクラスター \
--instance-groups ファイル://demo-cluster.json 

作成後、 amazon sagemaker describe-clusterおよびamazon sagemaker list-cluster-nodesを使用して、クラスターとノードの詳細を確認できます。コントローラーノードのクラスターIDとインスタンスIDをメモしておいてください。この情報はクラスターに接続するために必要です。

Amazon FSx for Lustre などの共有ファイルシステムをアタッチすることもできます。FSx for Lustre を使用するには、Amazon Virtual Private Cloud (Amazon VPC) 構成を使用してクラスターをセットアップする必要があります。これは、SageMaker VPC の作成方法と FSx for Lustre のデプロイ方法を示す Amazon CloudFormation テンプレートです。

クラスターを接続する

クラスターユーザーは、クラスター管理者からクラスターへのアクセス権限を付与される必要があります。アクセス権限が付与されると、SSH経由でクラスターに接続し、ジョブのスケジュール設定や実行が可能になります。Amazon Systems ManagerにプリインストールされているAmazon CLIプラグインを使用して、クラスターのコントローラーノードに接続できます。

デモでは、コントローラーノードのクラスター ID とインスタンス ID をターゲットとして指定して、次のコマンドを実行しました。

 AWS SSM 開始セッション
--target sagemaker-cluster:ntg44z9os8pn_controller-group-i-05a854e0d4358b59c \
--リージョン us-west-2 

Slurm を使用して、クラスター上でジョブをスケジュールおよび実行します。

SageMaker HyperPodはリリース以来、Slurmによるワークロードオーケストレーションをサポートしています。Slurmは、人気の高いオープンソースのクラスター管理およびジョブスケジューリングシステムです。クラスター作成時に、ライフサイクルスクリプトを使用してSlurmをインストールおよび設定できます。サンプルのライフサイクルスクリプトでそのプロセスを確認できます。その後、標準的なSlurmコマンドを使用してジョブをスケジュールし、開始できます。詳細なアーキテクチャ情報と便利なコマンドについては、Slurmクイックスタートユーザーガイドをご覧ください。

このデモでは、Amazon Web Services 機械学習トレーニングリファレンスアーキテクチャ GitHub リポジトリの例を使用しました。この例では、Trn1 インスタンスを使用して Slurm で Llama 2 7B をトレーニングする方法を示しています。私のクラスターは既に Slurm でセットアップされており、FSx for Lustre ファイルシステムがマウントされています。

知らせ

Llama 2モデルはMetaによって管理されています。Metaリクエストを通じてページへのアクセスをリクエストできます。

クラスタ環境の設定

SageMaker HyperPodは、Conda、venv、Docker、enrootなど、様々な環境でのトレーニングをサポートしています。READMEの指示に従って仮想環境amazon_neuron_venv_pytorchを構築し、torch_neuronxとneuronx-nemo-megatronライブラリをセットアップして、Trn1インスタンスでモデルをトレーニングしました。

モデル、単語セグメンター、データセットを準備します。

指示に従って、Llama 2モデルとトークナイザーをダウンロードし、モデルをHugging Face形式に変換しました。次に、RedPajamaデータセットをダウンロードし、トークン化しました。最後の準備として、Llama 2モデルをAOT（Ahead-of-Time）コンパイルで事前コンパイルし、モデルのトレーニングを高速化しました。

クラスター上でジョブを開始する

これで、 sbatchコマンドを使用してモデル トレーニング ジョブを開始する準備が整いました。

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

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

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

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

モデルのトレーニングが開始されたことを示す出力例を以下に示します。

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

モデルトレーニングジョブをさらに監視および分析するには、SageMaker を使用して TensorBoard またはその他のツールをホストできます。

発売中

SageMaker HyperPod は現在、米国東部 (オハイオ)、米国東部 (バージニア北部)、米国西部 (オレゴン)、アジアパシフィック (シンガポール)、アジアパシフィック (シドニー)、アジアパシフィック (東京)、欧州 (フランクフルト)、欧州 (アイルランド)、欧州 (ストックホルム) の Amazon Web Services リージョンでご利用いただけます。

もっと詳しく知る：

● 料金情報とサポートされているクラスターインスタンスタイプのリストについては、Amazon SageMaker HyperPod を参照してください。

● 開発者ガイドを見る

● Amazon Web Services コンソールにアクセスして、SageMaker HyperPod を使用して基本モデルのトレーニングを開始します。

出典: https://aws.amazon.com/blogs/aws/introducing-amazon-sagemaker-hyperpod-a-purpose-built-infrastructure-for-distributed-training-at-scale?trk=cndc-detail/

— アンチェ

追記：Amazon Web Services (AWS) に関する各記事には著者名が 1 人しか記載されていませんが、これらのブログ記事はチームの努力の成果です。サンプルコードの作成に多大なご協力をいただき、大規模モデル学習インフラストラクチャ、Slurm、SageMaker HyperPod の管理に関する専門知識を共有していただいた Brad Doran、Justin Pirtle、Ben Snyder、Pierre-Yves Aquilanti、Keita Watanabe、Verdi March の皆様に感謝申し上げます。