Trainium Slurm クラスターの事前トレーニングチュートリアル - Amazon SageMaker AI
Slurm クラスターで Trainium 環境を設定するTrainium でトレーニングジョブを起動する

Trainium Slurm クラスターの事前トレーニングチュートリアル

次のチュートリアルでは、Slurm クラスターに Trainium 環境を設定し、Llama 80 億パラメータモデルでトレーニングジョブを開始します。

前提条件

環境のセットアップを開始する前に、以下を確認します。

  • SageMaker HyperPod Trainium Slurm クラスターをセットアップします。

  • 共有ストレージの場所。クラスターノードからアクセスできる Amazon FSx ファイルシステムまたは NFS システムでかまいません。

  • 以下の形式のいずれか。

    • JSON

    • JSONGZ (圧縮 JSON)

    • ARROW

  • (オプション) HuggingFace のモデル重みを事前トレーニングまたはファインチューニングに使用する場合は、HuggingFace トークンを取得する必要があります。アクセストークンの詳細については、「ユーザーアクセストークン」を参照してください。

Slurm クラスターで Trainium 環境を設定する

Slurm クラスターでトレーニングジョブを開始するには、次の手順を実行します。

  • クラスターのヘッドノードに SSH 接続します。

  • ログインしたら、Neuron 環境を設定します。Neuron のセットアップについては、「Neuron のセットアップ手順」を参照してください。Neuron のドライバーがプリインストールされている Deep learning AMI (DLAMI Pytorch 搭載の Ubuntu 20 など) を使用することをお勧めします。

  • クラスター内の共有ストレージ場所にSageMaker HyperPod レシピリポジトリのクローンを作成します。共有ストレージの場所は、クラスターノードからアクセスできる Amazon FSx ファイルシステムまたは NFS システムでかまいません。

    git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
cd sagemaker-hyperpod-recipes
pip3 install -r requirements.txt

  • 次のチュートリアルを実行します。HuggingFace Llama3-8B 事前トレーニング

  • モデルの設定を準備します。Neuron リポジトリで使用できるモデル設定。このチュートリアルで使用するモデル設定については、「llama3 8b モデル設定」を参照してください。

Trainium でトレーニングジョブを起動する

Trainium でトレーニングジョブを起動するには、クラスター設定と Neuron レシピを指定します。例えば、Trainium で llama3 8b 事前トレーニングジョブを起動するには、起動スクリプト launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh を以下のとおり設定します。

  • MODEL_CONFIG: 環境設定セクションのモデル設定

  • (オプション) 次の key-value ペアを設定することで、HuggingFace から事前トレーニング済みの重みが必要な場合は、HuggingFace トークンを指定できます。

    recipes.model.hf_access_token=<your_hf_token>
#!/bin/bash

#Users should set up their cluster type in /recipes_collection/config.yaml

SAGEMAKER_TRAINING_LAUNCHER_DIR=${SAGEMAKER_TRAINING_LAUNCHER_DIR:-"$(pwd)"}

COMPILE=0
TRAIN_DIR="${TRAIN_DIR}" # Location of training dataset
MODEL_CONFIG="${MODEL_CONFIG}" # Location of config.json for the model

HYDRA_FULL_ERROR=1 python3 "${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py" \
    base_results_dir="${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results" \
    instance_type="trn1.32xlarge" \
    recipes.run.compile="$COMPILE" \
    recipes.run.name="hf-llama3-8b" \
    recipes.trainer.num_nodes=4 \
    recipes=training/llama/hf_llama3_8b_seq8k_trn1x4_pretrain \
    recipes.data.train_dir="$TRAIN_DIR" \
    recipes.model.model_config="$MODEL_CONFIG"

以下のコマンドを実行して、トレーニングジョブを起動します。

bash launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh

Slurm クラスターの設定に関する詳細は、「HyperPod Slurm でのトレーニングジョブの実行」を参照してください。