基于 Amazon SageMaker 优化 Alpaca-lora 模型

项目的基本情况

斯坦福大学开源的 stanford_alpaca 项目（https://github.com/tatsu-lab/stanford_alpaca），基于“指令遵循数据”对 META 开源的大语言模型 LLAMA 进行微调（supervised fine-tuning）。出色的问答表现，使它一经发布就迅速成为了类 ChatGPT 大模型研发的主要开源技术方案。然而为了在 stanford_alpaca 中对 LLAMA 进行微调，对机器的显卡和硬件性能有很高的要求：在实验中我们发现，即使采用 p4d.24xlarge 这样装有 8 张 A100 显卡（每张卡显存是 40G）的机器，也无法把 stanford_alpaca 的原版代码跑起来（除非采用 deepspeed 进行改造）。

因此，如果要采用 stanford_alpaca 对 70 亿参数的 LLAMA 进行微调，对于机器的配置要求较高，不利于在行业内大规模使用。为了解决这个痛点，开源项目 alpaca_lora（https://github.com/tloen/alpaca-lora）应运而生。alpaca_lora 的核心思路是通过 LoRA（low-rank adaptation）的技术，大幅度减少需要训练的参数量（trainable parameters），从而大大提高了模型训练速度，也降低了对于显卡配置的需求。经过 alpaca_lora 的作者实验，当训练数据有 5 万条时，单张 RTX 4090 的显卡只使用几个小时即可完成 alpaca_lora 的训练。

LoRA 技术的基本思想

LoRA 方法来源于论文《LORA: LOW-RANK ADAPTATION OF LARGE LAN- GUAGE MODELS》。

LoRA 提出，为了用最小的成本达到微调大模型的目的，可以冻结原始预训练模型的参数，仅更新训练新增的小量级参数。与原始大模型的千亿参数量（比如 GPT3的参数量有 175B）相比，LoRA 将需要训练的参数量减小了 10,000 倍。

可以用一张图来解释 LoRA 的思想：

这里：h = W₀ X+ ∆W X = W₀ X + B A X

左侧蓝色部分为预训练大模型（记为 W），在训练的过程中所有的参数都会被冻结，不进行任何更新。

右侧 A、B 两部分包含了可训练的参数。其中对矩阵 A 使用随机高斯分布进行初始化，对矩阵 B 使用 0 进行初始化。所以，∆W = BA 在最初对模型结果无影响。

LoRA 的主要优势有：

由于不需要对大部分参数进行更新，提升了训练的效率，降低了训练的硬件门槛。

用简明的线性结合，可以将训练的矩阵与原始冻结的权重相结合，与整个模型进行训练的情况相比，没有引入延迟。

Alpaca LoRA 项目的实现细节

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

首先我们来看 Dockerfile：

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker
WORKDIR /opt/ml
COPY requirements.txt /opt/ml/
RUN pip install -r requirements.txt 
COPY . /opt/ml
RUN ls /opt/ml
# Make all local GPUs visible
ENV NVIDIA_VISIBLE_DEVICES="all"
ENTRYPOINT ["/opt/ml/train"]

解读：

• FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker：这个命令指定了基础镜像，即所构建的镜像是在已有的镜像之上构建的。在这个命令中，基础镜像是一个由亚马逊云科技提供的 PyTorch 训练镜像，版本为 1.13.1，支持 GPU，使用 Python 3.9，CUDA 11.1，运行在 Ubuntu 20.04 操作系统上。
• WORKDIR /opt/ml：这个命令设置了当前工作目录为/opt/ml。这是一个惯例，/opt/ml 是 Amazon SageMaker 的默认工作目录。
• COPY requirements.txt /opt/ml/：这个命令将当前目录下的 requirements.txt 文件复制到/opt/ml 目录下。这个文件包含了所有需要安装的 Python 依赖项。
• RUN pip install -r requirements.txt：这个命令使用 pip 安装所有需要的 Python 依赖项，以确保在 Docker 镜像中存在这些依赖项。
• COPY . /opt/ml：这个命令将当前目录下的所有文件和文件夹复制到/opt/ml 目录下。
• RUN ls /opt/ml：这个命令列出/opt/ml 目录下的所有文件和文件夹。这个命令不是必须的，但是可以用来验证 COPY 命令是否成功将文件复制到/opt/ml 目录下。
• ENV NVIDIA_VISIBLE_DEVICES="all"：这个命令设置了环境变量。NVIDIA_VISIBLE_DEVICES 的值为”all”。这个环境变量的作用是将所有可用的本地 GPU 设备暴露给 Docker 镜像中运行的进程。
• ENTRYPOINT ["/opt/ml/train"]：这个命令设置了 Docker 镜像的默认入口点为/opt/ml/train。这个命令指定了在启动 Docker 容器时要运行的命令。在这个例子中，train 是一个可执行文件，用于启动 PyTorch 训练过程。

接下来我们看可执行文件 train：

#!/bin/bash
WORLD_SIZE=4 CUDA_VISIBLE_DEVICES=0,1,2,3 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=4 finetune.py --base_model  'David003/llama-7b-hf-20230407' --data_path 'test.json' --output_dir '/opt/ml/lora-alpaca-train'  --batch_size 128 --micro_batch_size 4

/opt/ml/s5cmd sync /opt/ml/lora-alpaca-train/ <s3_path>

解读：

• #!/bin/bash：这个命令指定了当前脚本的解释器，即使用 Bash 解释器来运行这个脚本。
• WORLD_SIZE=4 CUDA_VISIBLE_DEVICES=0,1,2,3 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=4 finetune.py --base_model 'David003/llama-7b-hf-20230407' --data_path 'test.json' --output_dir '/opt/ml/lora-alpaca-train' --batch_size 128 --micro_batch_size 4：这个命令使用 torchrun 来运行 finetune.py 脚本，并指定了一些环境变量。其中，WORLD_SIZE=4 表示在 4 个进程中进行分布式训练；CUDA_VISIBLE_DEVICES=0,1,2,3 表示仅使用 4 个 GPU 设备；TORCH_DISTRIBUTED_DEBUG=DETAIL 用于调试分布式训练问题。finetune.py 脚本是训练脚本的名称，其它参数包括预训练模型名称、训练数据路径、输出目录、批量大小等。
• /opt/ml/s5cmd sync /opt/ml/lora-alpaca-train/ <s3_path>：这个命令使用 s5cmd 将训练输出目录中的文件同步到 Amazon S3 存储桶中。s5cmd 是一个 Amazon S3 命令行工具，用于从 Amazon SageMaker 中上传或下载文件。在这个命令中，源目录是/opt/ml/lora-alpaca-train/，目标存储桶是<s3_path>。

最后我们看 main 脚本：

!sh build_push.sh lora

build_push.sh 用于 build 并 push 镜像到名为 lora 的 ECR 仓库中

from sagemaker import get_execution_role

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

estimator.fit(''s3://<data path> ')

解读：

• from sagemaker import get_execution_role：导入 Amazon SageMaker 的 Python SDK，并使用 get_execution_role()函数获取当前执行角色。
• role = get_execution_role()：将执行角色分配给变量 role。
• from sagemaker.estimator import Estimator：导入 Estimator 类，该类用于定义和运行 Amazon SageMaker 训练作业。
• instance_type = 'ml.g5.12xlarge'：定义实例类型，即用于训练作业的 EC2 实例类型。
• estimator = Estimator(role=role, instance_count=1, instance_type=instance_type, image_uri=<lora ECR路径>')：创建一个 Estimator 对象，指定执行角色、实例数量、实例类型和 Docker 镜像 URI。其中，image_uri 参数指定了 Docker 镜像 URI，该镜像包含了预安装的 PyTorch 和其他依赖项。
• estimator.fit('s3://<data path> ')：启动训练作业并开始训练。其中，fit()方法接收一个 S3 路径作为输入数据源。在训练过程中，Amazon SageMaker 会自动配置、启动和监视训练实例。

CloudWatch 记录

总结

Alpaca Lora 作者采用了 Hugging Face 的轻量化微调库（Parameter Efficient Fine-Tuning，PEFT）中所支持的 LoRA 方法。LoRA 方法的两项配置会直接影响需要训练的参数量：

1）LoRA 目标模块（lora_target_modules），用于指定要对哪些模块的参数进行微调。比如我们可以对 Q, K, V, O 都进行微调；也可以只对 Q、V 进行微调。不同的设定会影响需要微调的参数量，也会影响训练过程中的计算量。比如当我们设定只对 Q、V 进行微调时，需要训练的参数量（trainable parameters）只占整个模型参数总量的 6% 左右。

2）LoRA 的秩（lora_r）也是影响训练参数量的一个重要因素。客观来说，使用 LoRA 这样的方法训练得到的模型，在效果上必然会和直接在原始大模型基础上进行训练的效果有一定差异。因此，可以结合所拥有的机器配置、可以容忍的最大训练时长等因素，来灵活地配置 LoRA 的使用方法。

本篇作者