HyperPod Slurm 叢集預先訓練教學課程 (GPU) - Amazon SageMaker AI
HyperPod GPU Slurm 環境設定啟動訓練任務

HyperPod Slurm 叢集預先訓練教學課程 (GPU)

下列教學課程會設定 Slurm 環境，並在 Llama 80 億參數模型上啟動訓練任務。

先決條件

開始設定您的環境以執行配方之前，請確定您已：

  • 設定 HyperPod GPU Slurm 叢集。

    • 您的 HyperPod Slurm 叢集必須啟用 Nvidia Enroot 和 Pyxis (這些項目預設為啟用)。

  • 共用儲存位置。它可以是可從叢集節點存取的 Amazon FSx 檔案系統或 NFS 系統。

  • 採用下列其中一種格式的資料：

    • JSON

    • JSONGZ (壓縮 JSON)

    • ARROW

  • (選用) 如果您要使用來自 HuggingFace 的模型權重進行預先訓練或微調，則必須取得 HuggingFace 權杖。如需取得權杖的詳細資訊，請參閱使用者存取權杖。

HyperPod GPU Slurm 環境設定

若要在 HyperPod GPU Slurm 叢集上啟動訓練任務，請執行下列動作：

  1. 對 Slurm 叢集的主節點執行 SSH。

  2. 登入後，請設定虛擬環境。請確定您使用的是 Python 3.9 或更新版本。

    #set up a virtual environment
python3 -m venv ${PWD}/venv
source venv/bin/activate

  3. 將 SageMaker HyperPod 配方和 SageMaker HyperPod 轉接器儲存庫複製到共用儲存位置。

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

  4. 使用 Enroot 建立 squash 檔案。若要尋找 SMP 容器的最新版本，請參閱 SageMaker 模型平行化程式庫的版本備註。若要深入了解如何使用 Enroot 檔案，請參閱建置 AWS 最佳化的 Nemo-Launcher 映像。

    REGION="<region>"
IMAGE="658645717510.dkr.ecr.${REGION}.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121"
aws ecr get-login-password --region ${REGION} | docker login --username AWS --password-stdin 658645717510.dkr.ecr.${REGION}.amazonaws.com
enroot import -o $PWD/smdistributed-modelparallel.sqsh dockerd://${IMAGE}
mv $PWD/smdistributed-modelparallel.sqsh "/fsx/<any-path-in-the-shared-filesystem>"

  5. 若要使用 Enroot squash 檔案開始訓練，請使用下列範例來修改 recipes_collection/config.yaml 檔案。

    container: /fsx/path/to/your/smdistributed-modelparallel.sqsh

啟動訓練任務

在您安裝相依性之後，請從 sagemaker-hyperpod-recipes/launcher_scripts 目錄啟動訓練任務。您可以複製 SageMaker HyperPod 配方儲存庫來取得相依性：

首先，從 Github 挑選您的訓練配方，模型名稱會指定為配方的一部分。在以下範例中，我們使用 launcher_scripts/llama/run_hf_llama3_8b_seq16k_gpu_p5x16_pretrain.sh 指令碼搭配序列長度 8192 預先訓練配方 llama/hf_llama3_8b_seq16k_gpu_p5x16_pretrain 啟動 Llama 8b。

  • IMAGE：來自環境設定區段的容器。

  • (選用) 如果您需要來自 HuggingFace 的預先訓練權重，您可以設定下列金鑰/值對，以提供 HuggingFace 權杖：

    recipes.model.hf_access_token=<your_hf_token>
#!/bin/bash
IMAGE="${YOUR_IMAGE}"
SAGEMAKER_TRAINING_LAUNCHER_DIR="${SAGEMAKER_TRAINING_LAUNCHER_DIR:-${PWD}}"

TRAIN_DIR="${YOUR_TRAIN_DIR}" # Location of training dataset
VAL_DIR="${YOUR_VAL_DIR}" # Location of validation dataset

# experiment ouput directory
EXP_DIR="${YOUR_EXP_DIR}"

HYDRA_FULL_ERROR=1 python3 "${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py" \
  recipes=training/llama/hf_llama3_8b_seq16k_gpu_p5x16_pretrain \
  base_results_dir="${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results" \
  recipes.run.name="hf_llama3_8b" \
  recipes.exp_manager.exp_dir="$EXP_DIR" \
  recipes.model.data.train_dir="$TRAIN_DIR" \
  recipes.model.data.val_dir="$VAL_DIR" \
  container="${IMAGE}" \
  +cluster.container_mounts.0="/fsx:/fsx"

在啟動器指令碼中設定了所有必要參數之後，您可以使用下列命令執行指令碼。

bash launcher_scripts/llama/run_hf_llama3_8b_seq16k_gpu_p5x16_pretrain.sh

如需 Slurm 叢集組態的詳細資訊，請參閱 在 HyperPod Slurm 上執行訓練任務。