基于 Amazon SageMaker 优化 Stanford Alpaca 模型

在类 ChatGPT 大模型方面的开源项目中，Stanford Alpaca（https://github.com/tatsu-lab/stanford_alpaca）无疑是最受瞩目的项目之一。它的核心贡献有两个：

1）构造了一个生成“指令遵循数据”的 pipeline（ instruction-following data generation pipeline）

基本思路是构造一些 prompt 发送给 OpenAI 的大模型（如 text-davinci-003），在这些 prompt 中，告诉大模型需要它产生 20 个不同的任务指令（a set of 20 diverse task instructions），并提供 3 个种子任务示例（从 175 个 seed task 中每次随机选 3 个）。

如下是一个发给 text-davinci-003 的 prompt 的例子：

1. Instruction: Replace the <mask> token in the  text with proper words that are consistent with the context. You can use  multiple words for each <mask> token.
1. Input:
The Tom and Jerry episode The Cat Concerto, which  features <mask>, motivated two-year-old Lang Lang to learn the piano.
1. Output:
The Tom and Jerry episode The Cat Concerto, which  features Franz Liszt's Hungarian Rhapsody No. 2, motivated two-year-old Lang  Lang to learn the piano.

text-davinci-003 这样的大模型会根据任务指令的要求，模仿“种子任务示例”，写完剩下的 17 个任务。

通过这样的方式，作者一共采集到 52000 条“指令遵循数据”，即 alpaca_data.json。

2）提供了基于“指令遵循数据”对 LLAMA 进行微调（supervised fine-tuning）的代码

按照 OpenAI 官方介绍（ https://openai.com/blog/chatgpt），ChatGPT 的思路分为三个步骤：

借鉴 ChatGPT 的思路，研发一个类 ChatGPT 有三个步骤：

第一步：采集一定的标注数据（prompt+期望的回答），在一个大语言模型（例如 GPT3 GPT-3.5）基础上，进行有监督的训练，得到“模型 A”。

第二步：采集对比数据（给定一个 prompt 以后，第一步的模型会有多个输出，标注员给出排序；据此标注一定量的数据），训练得到 Reward 模型，简称为“模型 B”（“模型 B”的模型结构与“模型 A”不同）。

第三步：基于第一步、第二步的模型基于 PPO 强化学习算法，训练得到最终的模型，简称为“模型 C”（“模型 C”的模型结构与“模型 A”相同）。

在类 ChatGPT 大模型的研发过程中，为了进行第一步的训练，目前通常使用 OPT、BLOOM、GPT-J、LLAMA 等开源大模型替代 GPT3、GPT3.5 等模型。Stanford Alpaca 提供了基于“指令遵循数据”对 LLAMA 进行微调（supervised fine-tuning）的代码，完成了“类 ChatGPT 大模型训练步骤”中的第一步。

在本文中，我们探索如何在 SageMaker 进行 Alpaca supervised fine-tuning。在这篇 blog 中，我们将采用自建镜像（BYOC）的方式。本次代码来源于 https://github.com/tatsu-lab/stanford_alpaca。由于合规原因，我们不能提供原始代码与数据，但是可以基于 SageMaker 的运行流程进行 Dockerfile 与 Entrypoint 解读。

首先我们来看 Dockerfile:

From 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04 
ENV LANG=C.UTF-8
ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/ml:${PATH}"
WORKDIR /opt/ml
COPY . /opt/ml
RUN ls /opt/ml
RUN python3 -m pip install -r requirements.txt 
RUN cd transformers_src/transformers-llama_push/ && python3 -m pip  uninstall -y transformers && python3 setup.py install && cd ..
RUN pip3 uninstall -y deepspeed && pip3 install deepspeed
ARG SAGEMAKER_PROGRAM=/opt/ml/train.py
ENV SAGEMAKER_PROGRAM=${SAGEMAKER_PROGRAM}

# Make all local GPUs visible
ENV NVIDIA_VISIBLE_DEVICES="all"
ENTRYPOINT ["/opt/ml/train"]

这个 Dockerfile 的作用是构建一个基于 PyTorch 的 GPU 训练环境的 Docker 镜像。首先，从 Amazon Elastic Container Registry 中拉取了一个名为 huggingface-pytorch-training:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04 的基础镜像。这个基础镜像包含了 Ubuntu 20.04 操作系统、Python 3.9 和 PyTorch 训练环境，并且已经安装了 Hugging Face Transformers 库的 4.26.0 版本。

然后，通过 ENV 命令设置了一些环境变量，将 LANG 设置为 C.UTF-8，将 PYTHONUNBUFFERED 和 PYTHONDONTWRITEBYTECODE 设置为 TRUE，以优化 Python 的性能和行为。

接下来，使用 WORKDIR 命令将工作目录设置为 /opt/ml，并使用 COPY 命令将当前目录中的所有文件复制到 Docker 镜像中的 /opt/ml 目录中。

然后，运行了一些命令，包括列出 /opt/ml 目录中的文件，安装 requirements.txt 中列出的 Python 依赖包，卸载了名为 deepspeed 的 Python 包，并安装了一个新版本。

接着，它使用 ARG 和 ENV 命令设置了一个名为 SAGEMAKER_PROGRAM 的环境变量，并将其设置为 /opt/ml/train.py。最后，它使用 ENV 命令将环境变量 NVIDIA_VISIBLE_DEVICES 设置为 all，以便让所有本地 GPU 可用，并使用 ENTRYPOINT 命令设置了一个程序入口点 /opt/ml/train。

之后我们看 train 脚本：

#!/bin/bash
df -h
deepspeed --num_gpus=8 train.py --deepspeed ds.json --model_name_or_path "decapoda-research/llama-7b-hf" --data_path alpaca_data.json     --output_dir "/tmp/llama_out"     --num_train_epochs 1     --per_device_train_batch_size 2     --per_device_eval_batch_size  1     --gradient_accumulation_steps 2     --evaluation_strategy "no"     --save_strategy "steps"     --save_steps 5     --save_total_limit 3     --learning_rate 2e-5     --weight_decay 0.     --warmup_ratio 0.03     --lr_scheduler_type "cosine"     --logging_steps 1     --cache_dir '/tmp' --fp16_full_eval     --fp16

s5cmd sync /tmp/llama_out s3://sagemaker-us-east-1-652582962535/llama/

使用 deepspeed 命令调用 train.py 脚本，并指定了一系列参数，包括使用 8 个 GPU 进行训练，模型名称或路径为 decapoda-research/llama-7b-hf，数据路径为 alpaca_data.json，输出目录为 /tmp/llama_out，训练时长为 1 个 epoch，每个训练批次的大小为 2，每个评估批次的大小为 1，梯度累积步数为 2，评估策略为“不评估”，保存策略为“按步保存”，每隔 5 个步骤保存一次，最多保存 3 次，学习率为 2e-5，权重衰减为 0，预热比例为 0.03，学习率调度类型为“余弦退火”，记录日志的步骤为 1，缓存目录为 /tmp，全精度评估使用 FP16，同时使用混合精度训练。

最后，使用 s5cmd 命令将 /tmp/llama_out 目录中的文件同步到名为 sagemaker-us-east-1-652582962535/llama 的 S3 存储桶中。

最后封装 estimator 并进行 fit 操作

from sagemaker import get_execution_role

role = get_execution_role()
from sagemaker.estimator import Estimator
instance_type = 'ml.p4d.24xlarge'
estimator = Estimator(role=role,
                      instance_count=1,
                      instance_type=instance_type,
                      image_uri='<ECR路径>')

estimator.fit('<S3路径>')

CloudWatch 记录

除了 Llama 外，其他自定义环境如 Vicuna 也可以通过同样的流程进行整理，测试与部署。

亚马逊AWS官方博客

基于 Amazon SageMaker 优化 Stanford Alpaca 模型

本篇作者