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 토큰을 받아야 합니다. 토큰 가져오기에 대한 자세한 내용은 User access tokens 섹션을 참조하세요.

Slurm 클러스터에서 Trainium 환경 설정

Slurm 클러스터에서 훈련 작업을 시작하려면 다음을 수행합니다.

  • Slurm 클러스터의 헤드 노드에 SSH로 연결합니다.

  • 로그인한 후 Neuron 환경을 설정합니다. Neuron 설정에 대한 자세한 내용은 Neuron setup steps을 참조하세요. Ubuntu 20 with DLAMI Pytorch와 같은 Neuron 드라이버가 사전 설치된 딥 러닝 AMI를 사용하는 것이 좋습니다.

  • 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 Pretraining 자습서를 살펴봅니다.

  • 모델 구성을 준비합니다. Neuron 리포지토리에서 사용할 수 있는 모델 구성입니다. 이 자습서에서 사용된 모델 구성은 llama3 8b model config를 참조하세요.

Trainium에서 훈련 작업 시작

Trainium에서 훈련 작업을 시작하려면 클러스터 구성과 Neuron 레시피를 지정합니다. 예를 들어 Trainium에서 llama3 8b 사전 훈련 작업을 시작하려면 시작 스크립트인 launcher_scripts/llama/run_hf_llama3_8b_seq8k_trn1x4_pretrain.sh를 다음과 같이 설정합니다.

  • MODEL_CONFIG: 환경 설정 섹션의 모델 구성

  • (선택 사항) 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에서 훈련 작업 실행 섹션을 참조하세요.