在 Amazon EKS 上使用 AWS Trainium 训练 Llama2

导言

生成式人工智能不仅改变着企业的运作方式，也加快了在更广泛的人工智能领域中创新的步伐。这股变革力量正在重新定义企业使用人工智能的方式。这些曾经一度被认为遥不可及的技术，目前已经在创造类人文本、图像、代码和音频等场景中被投入使用。生成式人工智能提供了一系列应用，而不仅仅是执行提示和促进互动对话。这些大语言模型越来越多地被应用于代码生成、内容摘要、数据分析等各种场景。被各行各业大规模地采用，彰显了大语言模型在解决众多复杂任务和挑战方面的实用性和通用性。然而，这些优势也带来了一系列新的挑战，尤其是在训练和运行这些大型模型方面。

大语言模型和生成式人工智能的规模不断扩大，大大增加了计算需求，导致开发和部署成本不断上升。随着数据规模和模型复杂度的不断增加，训练模型所需的资源也越来越多。这种趋势凸显了类似 Amazon EKS 这类高性价比解决方案的重要性。Amazon EKS 提供了必需的可扩展性和计算能力，可在不产生高昂费用的情况下灵活管理模型训练的工作负载。根据 TIRIAS Research的预测，到 2028 年，人工智能基础设施成本将超过 760 亿美元。现有的商业框架难以将这些不断增长的成本转嫁给消费者，因此必须采用新的商业模式或大幅降低成本，以确保人工智能的持续增长和消费者的持续受益。

在成本不断上涨和全球计算资源供应日益稀缺的情况下，AWS Trainium 为面临这些挑战的开发人员提供了一个实用的解决方案。通过使用 AWS Trainium，开发人员可以将训练模型的成本最多降低 50%，同时还能优化分布式训练的性能。这使 AWS Trainium 受到了那些希望在深度学习和模型开发领域中降本增效的人的追捧。想了解更多有关 AWS Trainium 的详细信息，请访问 AWS Trainium 产品页面。

使用 Amazon EKS 和 AWS Trainium 的分布式训练架构

该解决方案基于 Data on EKS，通过 Data on EKS，用户可快速简易地创建一个包含 Trn1 EC2 实例的节点组的 Amazon EKS 集群。每一个大小为 trn1.32xlarge 的 EC2 实例都有16 个 AWS Trainium 芯片，可用于可扩展、高性能和高成本效益的模型训练。在节点组中，所有的 Trn1 EC2 实例都是通过高速、低延迟的 EFA（Elastic Fabric Adapter）网络连接，以实现分布式训练所要求的集体通信。

每个 Llama 训练作业都是通过一个 Kubernetes pod 执行的，其中用的容器镜像包含了 Neuron SDK （Trn1 类型实例所需）和 AWS Neuron Reference for NeMo Megatron。这些软件提供了先进的训练策略和功能，包括数据、张量、管道和序列并行、selective activation checkpointing 和 ZeRO-1 optimizer sharding。

Kubernetes MPI Operator 用于协调多个 pod 上的训练作业，其中每个 worker pod 都运行在单个大小为 trn1.32xlarge 的 EC2 实例上。

Amazon FSx for Lustre 用作共享存储系统。每一个 worker pod 都连接到该文件系统，用于存储数据集、tokenizer 文件、Llama 训练脚本、训练日志、编译制品和模型检查点。

解决方案概述

在 Amazon EKS 上使用 AWS Trainium 训练 Llama2

注：本文使用的是 Meta 公司的 Llama tokenizer，该 tokenizer 受用户许可保护，在下载 tokenizer 文件之前必须接受用户许可。请在此处申请访问权限，以确保您有访问 Llama 文件的权限。

先决条件

EC2 实例或者 Cloud9 实例，请确保它们有至少 100 GB 的磁盘空间
AWS CLI v2
kubectl
Git
Docker
Terraform
Python, pip, jq, unzip

如果需要在 EC2 实例上安装上述所有的软件，可以直接执行这个脚本（请记得在运行该脚本后退出并重新登录当前用户，以确保所有更改，尤其是 docker 组的更改，都会在当前用户中生效）。

第一步：克隆 Data on EKS 仓库并进入 trainium-inferentia 目录。

git clone https://github.com/awslabs/data-on-eks.git
cd data-on-eks/ai-ml/trainium-inferentia

MPI Operator 默认是不安装的，如下所示。

在本文中，我们将使用下列命令来设置执行环境的环境变量。

export TF_VAR_enable_mpi_operator=true
export TF_VAR_region=us-west-2
export TF_VAR_trn1_32xl_min_size=4
export TF_VAR_trn1_32xl_desired_size=4

注：截止至 2024 年 1 月 1 日，AWS Trainium 实例仅在 us-west-2，us-east-1 和 us-east-2 等区域可用。

第二步：执行安装脚本配置 Amazon EKS 集群以及所有需要的插件。

./install.sh

第三步：执行以下命令连接 Amazon EKS 集群。

aws eks update-kubeconfig --region us-west-2 --name trainium-inferentia

第四步：进入 examples/llama2 目录并执行脚本 1-llama2-neuronx-pretrain-build-image.sh 构建 neuronx-nemo-megatron 容器镜像并推送到 Amazon ECR 仓库。

cd examples/llama2/
./1-llama2-neuronx-pretrain-build-image.sh

提示输入区域时，请输入你的 Amazon EKS 集群所在的区域，如下所示。

整个镜像构建过程大概需要 10 分钟。

第五步：启动一个 CLI pod 并以此作为进入 FSx 存储的入口。

./2-launch-cmd-shell-pod.sh

执行以下命令检查 CLI pod 是否达到 “Running” 的状态。

kubectl get pod -w

第六步：使用以下命令连接 “Running” 状态的 CLI pod。

kubectl exec -it cli-cmd-shell -- /bin/bash

我们将通过 CLI pod 下载 Llama tokenizer 文件。首先，执行以下命令登录 huggingface。

huggingface-cli login

提示输入 token 时，我们可以从 Hugging Face 的网站中找到我们的 access token，如下所示。

第七步：执行以下 Python 代码下载 Llama2-7b 的 tokenizer 文件并保存到/shared/llama7b_tokenizer 目录中。

python3 <<EOF
import transformers
tok = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tok.save_pretrained("/shared/llama7b_tokenizer")
EOF

第八步：下载 RedPajama-Data-1T-Sample 数据集。

cd /shared
git clone https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T-Sample  data/RedPajama-Data-1T-Sample

第九步：使用包含 neuronx-nemo-megatron 的 preprocessing 脚本对数据集进行 tokenize。这一过程在大小为 trn1.32xlarge 的 EC2 实例上运行大概需要 60 分钟。

cd /shared

# Clone the neuronx-nemo-megatron repo, which includes the required scripts
git clone https://github.com/aws-neuron/neuronx-nemo-megatron.git

# Combine the separate redpajama files to a single jsonl file
cat /shared/data/RedPajama-Data-1T-Sample/*.jsonl > /shared/redpajama_sample.jsonl

# Run preprocessing script using llama tokenizer
python3 neuronx-nemo-megatron/nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py \
    --input=/shared/redpajama_sample.jsonl \
    --json-keys=text \
    --tokenizer-library=huggingface \
    --tokenizer-type=/shared/llama7b_tokenizer \
    --dataset-impl=mmap \
    --output-prefix=/shared/data/redpajama_sample \
    --append-eod \
    --need-pad-id \
    --workers=32

注：当我们稍后在 Amazon EKS 中启动训练作业时，worker pod 将从 FSx 上 neuronx-nemo-megatron/nemo/examples 目录中运行训练脚本。这样做的便捷性在于，您可以直接在 Amazon FSx 上修改训练脚本，而无需每次因为更改训练脚本都需要重建 neuronx-nemo-megatron 容器。

第十步：执行以下命令更新 test_llama.sh 脚本。更新的目的在于为 worker pod 指定 Llama tokenizer 和数据集在 Amazon FSx 文件系统上的位置。

sed -i 's#^\(: ${TOKENIZER_PATH=\).*#\1/shared/llama7b_tokenizer}#' /shared/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/test_llama.sh
sed -i 's#^\(: ${DATASET_PATH=\).*#\1/shared/data/redpajama_sample_text_document}#' /shared/neuronx-nemo-megatron/nemo/examples/nlp/language_modeling/test_llama.sh

更改之前：

更改之后：

第十一步：执行以下命令清理 CLI pod。

kubectl delete pod cli-cmd-shell

第十二步：现在我们可以开始预编译和训练作业了！在执行训练任务之前，我们首先需要执行预编译任务，以便准备模型制品。该步骤提取并编译 Llama2-7B 模型的底层计算图，并生成可在 AWS Trainium 芯片上运行的 AWS Neuron 可执行文件（NEFF）。这些 NEFF 文件保存在位于 Amazon FSx 上的 AWS Neuron 持久缓存中，以便训练作业读取它们。

在运行编译任务前，请执行此命令以确保 MPI Operator 正常工作。

kubectl get all -n mpi-operator

通过以下命令执行预编译脚本。

./3-llama2-neuronx-mpi-compile.sh

整个预编译过程大概需要 10 分钟。期间可以通过执行以下命令检查编译任务是否已经完成。

kubectl get po | grep compile

第十三步：当预编译任务完成后，执行以下脚本启动在 4 个大小为 trn1.32xlarge 的节点上运行的训练任务。

./4-llama2-neuronx-mpi-train.sh

第十四步：获取训练任务的 UID。首先，执行以下命令找到训练任务对于的 pod。

kubectl get pods | grep launcher

确认训练的 pod 是 “Running” 状态之后，执行以下命令找到对应的 UID（注意将 test-mpi-train-launcher-xxx 换成您的 pod 的名字）。

kubectl get pod test-mpi-train-launcher-g52f4 -o json | jq -r ".metadata.uid"

第十五步：利用上一步获取的 UID 查看训练日志。

kubectl exec -it test-mpi-train-worker-0 -- tail -f /shared/nemo_experiments/<UID>/0/log

第十六步：使用 neuron-top 命令监控 AWS Trainium 芯片的使用率。

kubectl exec -it test-mpi-train-worker-0 -- /bin/bash -l neuron-top

第十七步：部署可视化工具 Tensorboard ，用以将训练过程的各项指标可视化展示。

./5-deploy-tensorboard.sh

脚本执行完成后会输出一个 Loadbalancer 地址，在浏览器中访问该地址即可看到 TensorBoard 界面，如下所示。

清理资源

实验完成后，可执行以下清理脚本删除上述部署的所有资源。

cd data-on-eks/ai-ml/trainium-inferentia
./cleanup.sh

结论

在本文中，我们向您展示了 AWS Trainium 如何与 Amazon EKS 上的 Neuronx-nemo- megatron 集成。这标志着我们在应对人工智能模型训练中不断增长的计算需求和成本挑战方面取得了重大的进展。值得注意的是，AWS Trainium 可节省高达 50% 的训练成本，同时还具有高性能的特点。再加上 Neuron SDK 与流行的机器学习框架的兼容性，为人工智能模型训练提供了绝佳的环境。MPI Operator 和 Data on Amazon EKS（DoEKS）更进一步提高了分布式训练的效率和可扩展性。ZeRO-1 optimizer sharding 和 selective activation checkpointing 等创新功能不仅使最前沿的机器学习研究的门槛降低了，同时还推动了人工智能行业实现前所未有的创新。