在 SageMaker Hyperpod 上使用預構建容器對 Llama 3.2 1B 進行持續預訓練

本教程演示瞭如何使用 Hugging Face Optimum Neuron 庫在 Amazon SageMaker Hyperpod 上持續預訓練 Llama 3.2 1B 模型。我們利用了張量並行、序列並行和 ZeRO-1 等多種效能最佳化技術，在由 Trainium 驅動的例項上高效地訓練大型語言模型。

使用 SageMaker Hyperpod 的一個主要好處是能夠利用 Hugging Face 提供的預構建 Optimum Neuron 容器。這些容器預裝了所有必需的庫和依賴項，使得在 AWS Trainium 例項上開始訓練變得非常容易。

透過使用 SageMaker 預構建的容器，您可以避免手動設定環境的麻煩，專注於核心的訓練和微調任務。這些容器針對性能進行了最佳化，幷包含了各種最佳化技術，例如張量並行和選擇性檢查點，以高效地訓練像 Llama 3.2 1B 這樣的大型語言模型。

您將學習如何

在 SageMaker Hyperpod 上使用預構建容器對 Llama 3.2 1B 進行持續預訓練

1. 設定 AWS 環境

在開始本教程之前，您需要設定您的 AWS 環境。

建立一個包含至少一個 `trn1.32xlarge` 例項的 AWS SageMaker Hyperpod 叢集。您可以按照 Hyperpod EKS 研討會的步驟來設定叢集。
由於 Llama 3.2 是一個受限訪問模型，使用者在執行此示例前必須在 Hugging Face 註冊並獲取一個訪問令牌。您還需要在 meta-llama/Llama-3.2-1B 模型頁面上審閱並接受許可協議。
配置您的 AWS 憑證。如果您還沒有設定 AWS 憑證，可以透過安裝 AWS CLI 並執行 aws configure 來完成。您需要輸入您的 AWS Access Key ID、Secret Access Key、預設區域和輸出格式。
```
aws configure
AWS Access Key ID [None]: YOUR_ACCESS_KEY
AWS Secret Access Key [None]: YOUR_SECRET_KEY
Default region name [None]: YOUR_REGION
Default output format [None]: json
```

2. 準備訓練環境

使用必要的依賴項設定您的訓練環境。

git clone https://github.com/huggingface/optimum-neuron.git
mkdir ~/pre-training
cd pre-training

cp -r ../optimum-neuron/docs/source/training_tutorials/amazon_eks .
cd amazon_eks

登入到 ECR 並拉取 `huggingface-pytorch-training-neuronx` 映象。

region=us-east-1
dlc_account_id=************
aws ecr get-login-password --region $region | docker login --username AWS --password-stdin $dlc_account_id.dkr.ecr.$region.amazonaws.com

docker pull ${dlc_account_id}.dkr.ecr.${region}.amazonaws.com/huggingface-pytorch-training-neuronx:2.1.2-transformers4.43.2-neuronx-py310-sdk2.20.0-ubuntu20.04-v1.0

構建 Docker 映象並將其推送到您的 ECR 登錄檔。

export AWS_REGION=$(aws ec2 describe-availability-zones --output text --query 'AvailabilityZones[0].[RegionName]')
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
export REGISTRY=${ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com/
export IMAGE=optimum-neuron-llama-pretraining
export TAG=:latest

docker build -t ${REGISTRY}${IMAGE}${TAG} .

將映象推送到您的私有登錄檔。

# Create registry if needed
export REGISTRY_COUNT=$(aws ecr describe-repositories | grep \"${IMAGE}\" | wc -l)
if [ "${REGISTRY_COUNT//[!0-9]/}" == "0" ]; then
   echo "Creating repository ${REGISTRY}${IMAGE} ..."
   aws ecr create-repository --repository-name ${IMAGE}
else
   echo "Repository ${REGISTRY}${IMAGE} already exists"
fi

# Login to registry
echo "Logging in to $REGISTRY ..."
aws ecr get-login-password | docker login --username AWS --password-stdin $REGISTRY

# Push image to registry
docker image push ${REGISTRY}${IMAGE}${TAG}

3. 配置訓練作業

接下來，您將生成預訓練作業要使用的指令碼。首先，使用您在先決條件步驟中提到的訪問令牌登入 Hugging Face。在執行之前，修改 `generate-jobspec.sh` 指令碼以包含 Hugging Face 訪問令牌。

export HF_ACCESS_TOKEN="<your_HF_token_here>"

透過執行 `generate-jobspec.sh` 生成 Kubernetes 作業規範。這將為 Amazon SageMaker Hyperpod EKS 叢集建立一個名為 `llama_train.yaml` 的部署清單。

./generate-jobspec.sh

4. 在 SageMaker Hyperpod 上啟動訓練

將訓練作業部署到您的 Kubernetes 叢集。

kubectl apply -f llama_train.yaml

該清單使用 torchrun 在叢集上執行訓練指令碼以進行分散式訓練。您可以在 run_clm.py 檢視完整的訓練指令碼。

您將在此指令碼中使用以下分散式訓練技術：

分散式訓練：每個節點使用 8 個程序的 torchrun，以實現高效的多裝置訓練。
模型並行：同時實現張量並行 (TP=8) 和流水線並行 (PP=1)。
混合精度：利用 BFloat16 提高訓練效率。
梯度檢查點：實現記憶體高效的訓練。

該清單在叢集上執行以下命令。環境變數是在 `generate-jobspec.sh` 中建立清單時設定的。

torchrun --nproc_per_node=8 --nnodes=${NUM_NODES} run_clm.py \
    --model_name_or_path=${HF_MODEL_NAME}
    --token=${HF_ACCESS_TOKEN}
    --dataset_name=${DATASET_NAME}
    --dataset_config_name=${DATASET_CONFIG_NAME}
    --streaming=True
    --cache_dir=${TOKENIZED_DATA_PATH}
    --num_train_epochs=1
    --do_train
    --learning_rate=1e-4
    --max_steps=${MAX_STEPS}
    --per_device_train_batch_size=${BATCH_SIZE}
    --per_device_eval_batch_size=4
    --gradient_accumulation_steps=1
    --gradient_checkpointing
    --block_size=4096
    --bf16
    --max_grad_norm=1.0
    --lr_scheduler_type=linear
    --tensor_parallel_size=8
    --pipeline_parallel_size=1
    --logging_steps=1
    --save_total_limit=1
    --output_dir=${CHECKPOINT_DIR}
    --overwrite_output_dir

訓練作業現在將開始在 SageMaker Hyperpod 叢集上執行。

這使用了來自 Optimum-neuron 的預構建指令碼。該指令碼使用 Optimum Neuron 庫中的 Trainer 類，這是 Hugging Face Trainer 的一個專門版本，專為在 AWS Trainium 例項上訓練而最佳化。

以下是指令碼中主要元件的概述：

模型載入：使用 `AutoModelForCausalLM.from_pretrained()` 載入模型，並採用延遲載入以實現並行處理。
資料處理：將資料集分詞並處理成適合語言建模的塊。
訓練引數：指令碼使用 `NeuronTrainingArguments` 來配置訓練超引數，包括張量並行和流水線並行的選項。
Trainer 設定：使用模型、訓練引數、資料集和其他必要元件建立一個 `[optimum.neuron.NeuronTrainer]` 的 Trainer 例項。
訓練迴圈：呼叫 `trainer.train()` 方法以開始持續的預訓練過程。

5. 監控和驗證訓練

您可以透過 Kubernetes 日誌監控進度。

# Monitor training logs
kubectl logs -f -n kubeflow llama-training-eks-worker-0

# Validate saved checkpoints
kubectl exec -it llama-training-eks-worker-0 -- ls -l /fsx/output