在 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