使用 Ray 和 PyTorch Lightning 在 AWS Trainium 上分布式微调 Llama 3.1 8B#

Run on Anyscale

本示例演示了如何使用 Ray Train、PyTorch Lightning 和 AWS Neuron SDK 在 AWS Trainium 实例上微调 Llama 3.1 8B 模型。

AWS Trainium 是 AWS 构建的机器学习 (ML) 芯片，用于 100B+ 参数模型的深度学习 (DL) 训练。AWS Neuron SDK 帮助开发者在 Trainium 加速器上训练模型。

准备环境#

有关设置利用 AWS Trainium 实例的 Amazon EKS 集群，请参阅设置 EKS 集群和工具。

创建 Docker 镜像#

EKS 集群准备就绪后，创建一个 Amazon ECR 仓库，用于构建和上传包含微调 Llama3.1 8B 模型所需 artifacts 的 Docker 镜像。

  1. 克隆 repo。

git clone https://github.com/aws-neuron/aws-neuron-eks-samples.git

  1. 进入 llama3.1_8B_finetune_ray_ptl_neuron 目录。

cd aws-neuron-eks-samples/llama3.1_8B_finetune_ray_ptl_neuron

  1. 触发脚本。

chmod +x 0-kuberay-trn1-llama3-finetune-build-image.sh
./0-kuberay-trn1-llama3-finetune-build-image.sh

  1. 输入集群运行的区域，例如：us-east-2。

  2. 在 AWS 控制台中验证 Amazon ECR 服务是否包含新创建的 kuberay_trn1_llama3.1_pytorch2 仓库。

  3. 更新用于创建 Ray 集群的 manifest 文件中的 ECR 镜像 ARN。

使用以下命令替换 1-llama3-finetune-trn1-create-raycluster.yaml 文件中的 和 占位符为实际值，以反映上面创建的 ECR 镜像 ARN。

export AWS_ACCOUNT_ID=<enter_your_aws_account_id> # for ex: 111222333444
export REGION=<enter_your_aws_region> # for ex: us-east-2
sed -i "s/<AWS_ACCOUNT_ID>/$AWS_ACCOUNT_ID/g" 1-llama3-finetune-trn1-create-raycluster.yaml
sed -i "s/<REGION>/$REGION/g" 1-llama3-finetune-trn1-create-raycluster.yaml

配置 Ray 集群#

AWS Neuron 示例仓库中的 llama3.1_8B_finetune_ray_ptl_neuron 目录简化了 Ray 配置。KubeRay 提供了一个 manifest，你可以将其应用到集群上以设置 head 和 worker Pod。

运行以下命令设置 Ray 集群：

kubectl apply -f 1-llama3-finetune-trn1-create-raycluster.yaml

访问 Ray Dashboard#

从集群进行端口转发以查看 Ray Dashboard 的状态，然后通过 http://localhost:8265 访问。使用以下命令在后台运行：

kubectl port-forward service/kuberay-trn1-head-svc 8265:8265 &

启动 Ray Jobs#

Ray 集群现在可以处理工作负载了。启动数据准备和微调的 Ray job。

  1. 启动 Ray job 下载 dolly-15k 数据集和 Llama3.1 8B 模型 artifacts。

kubectl apply -f 2-llama3-finetune-trn1-rayjob-create-data.yaml

  1. Job 成功执行后，运行以下微调 job：

kubectl apply -f 3-llama3-finetune-trn1-rayjob-submit-finetuning-job.yaml

  1. 通过 Ray Dashboard 监控 jobs。

有关以上每个步骤的详细信息，请参阅AWS 文档链接。