使用 Amazon SageMaker 微调 Baichuan-2 模型

本篇文章主要介绍如何使用 Amazon SageMaker 进行 Baichuan-2 模型微调的示例。

这个示例主要包括:

Baichuan-2 总体介绍
Baichuan-2 微调介绍
Baichuan-2 环境设置
Baichuan-2 微调训练

前言

随着 ChatGPT 的腾空出世, 国内外各种基座大语言竞相出炉，在其基础上衍生出种类繁多的应用场景。训练优异的基座大语言模型在通用性方面表现较好，但模型可能并未涉及到特定领域的专业术语、领域内的特定用语或上下文等。采用微调技术可以通过在领域特定数据上进行训练，使模型更好地适应目标领域的特殊语言模式和结构; 结合基座模型的通用性和领域特定性，使得模型更具实际应用价值。

Baichuan-2 总体介绍

Baichuan-2 是百川智能推出的新一代开源大语言模型，采用 2.6 万亿 Tokens 的高质量语料训练，新系列发布包含有 7B、13B 的 Base 和 Chat 版本，不仅继承了上一代良好的生成与创作能力，流畅的多轮对话能力以及部署门槛较低等众多特性，而且在数学、代码、安全、逻辑推理、语义理解等能力有显著提升。

Baichuan-2 微调介绍

模型微调主要分为 Full Fine-Tune 和 PEFT（Performance-Efficient Fine-Tune），前者模型全部参数都会进行更新，训练时间较长，训练资源较大; 而后者会冻结大部分参数、微调训练网络结构，常见的方式是 LoRA 和 P-Tuning v2。

PEFT 微调方式由于参数更新较少，可能导致模型无法学习到全部领域知识，对于特定任务或领域来说会出现推理不稳定的情况，因此大多数生产系统均使用全参数方式进行模型的微调。基于上述原因，本文会以全参数微调方式介绍 Baichuan-2 在 SageMaker 上的微调。

Baichuan-2 环境设置

备注: 项目中的示例代码均保存于代码仓库，地址如下: https://github.com/aws-samples/llm-workshop-on-amazon-sagemaker

升级 Python SDK
```
pip install -U sagemaker
```

获取运行时资源，包括区域、角色、账号、S3 桶等

import boto3
import sagemaker
from sagemaker import get_execution_role


sess                     = sagemaker.Session()
role                     = get_execution_role()
sagemaker_default_bucket = sess.default_bucket()

account                  = sess.boto_session.client("sts").get_caller_identity()["Account"]
region                   = sess.boto_session.region_name

Baichuan-2 微调训练

微调准备

克隆代码

# The version is 2023.10.18
git clone https://github.com/baichuan-inc/Baichuan2.git
cd Baichuan2
git reset --hard dde8bdb353bead91361e0e3e41df2bf3dc6d314e

下载 Baichuan-2 原始模型

from huggingface_hub import snapshot_download
from pathlib import Path


local_cache_path = Path("./model")
local_cache_path.mkdir(exist_ok=True)

model_name = "baichuan-inc/Baichuan2-13B-Chat"

# Only download pytorch checkpoint files
allow_patterns = ["*.json", "*.pt", "*.bin", "*.model", "*.py", "*.txt"]

# 2023.10.18
model_download_path = snapshot_download(
    repo_id=model_name,
    cache_dir=local_cache_path,
    allow_patterns=allow_patterns,
    revision='3b930dc387bf01202382cfe47742c988584d7876'
)

# Get the model files path
import os
from glob import glob

local_model_path = None

paths = os.walk(r'./model')
for root, dirs, files in paths:
    for file in files:
        if file == 'config.json':
            # print(os.path.join(root, file))
            local_model_path = str(os.path.join(root, file))[0:-11]
            print(local_model_path)
if local_model_path == None:
    print("Model download may failed, please check prior step!")

拷贝模型和数据到 S3

%%script env sagemaker_default_bucket=$sagemaker_default_bucket local_model_path=$local_model_path bash

chmod +x ./s5cmd
./s5cmd sync ${local_model_path} s3://${sagemaker_default_bucket}/llm/models/baichuan2/baichuan-inc/Baichuan2-13B-Chat/

rm -rf model

模型微调

模型的微调使用全参数模型，以实现微调后模型的稳定性
模型的微调使用开源框架 DeepSpeed 进行加速

准备基础镜像

使用 SageMaker 定制的深度学习训练镜像作为基础镜像，再安装 Baichuan-2 训练所需的依赖包。Dockerfile 如下：

## You should change below region code to the region you used, here sample is use us-west-2
From 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04

ENV LANG=C.UTF-8
ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE

RUN pip3 install deepspeed==0.10.0 \
    && pip3 install transformers==4.29.2 \
    && pip3 install jsonlines==3.1.0 \
    && pip3 install accelerate==0.22.0 \
    && pip3 install xformers==0.0.21 \
    && pip3 install peft==0.5.0

COPY src/Baichuan2/fine-tune/requirements.txt .
RUN pip3 install -r requirements.txt

## Make all local GPUs visible
ENV NVIDIA_VISIBLE_DEVICES="all"

模型微调代码

模型微调源代码较多，细节可以参考上述 git 仓库。下面主要重点强调 Baichuan2-13B-Chat 模型的特性:

单轮对话使用如下提示词：{‘role’: ‘user’, ‘content’: ‘咨询问题’}
多轮对话使用如下提示词：{‘role’: ‘user’, ‘content’: ‘咨询问题1’}, {‘role’: ‘assistant’, ‘content’: ‘模型回答’}, {‘role’: ‘user’, ‘content’: ‘咨询问题2’}
Baichuan2 有特定 Token 用来标识“角色”，例如 token 195 表示“用户角色”、token 196 表示“系统角色”等。

微调参数

为了节省显存，采用 DeepSpeed Stage-3
训练过程开启 bf16，实现整数范围和精度的平衡

数据集采用官方提供的 belle_chat_ramdon_10k.json

DEEPSPEED_OPTS="""
    Baichuan2/fine-tune/fine-tune.py 
    --deepspeed Baichuan2/fine-tune/ds_config.json 
    --model_name_or_path "/tmp/baichuan2/" 
    --data_path "Baichuan2/fine-tune/data/belle_chat_ramdon_10k.json" 
    --output_dir "/tmp/baichuan2_out" 
    --num_train_epochs 1 
    --per_device_train_batch_size 1 
    --per_device_eval_batch_size  1 
    --gradient_accumulation_steps 1 
    --evaluation_strategy "no" 
    --save_strategy "steps" 
    --save_steps 2000 
    --save_total_limit 2 
    --learning_rate 2e-5 
    --adam_beta1 0.9 
    --adam_beta2 0.98 
    --adam_epsilon 1e-8 
    --weight_decay 1e-4 
    --max_grad_norm 1.0 
    --warmup_ratio 0.0 
    --lr_scheduler_type "constant" 
    --logging_steps 1 
    --cache_dir '/tmp' 
    --model_max_length 512 
    --gradient_checkpointing True 
    --bf16 True 
    --tf32 True 
    --report_to "none"
"""

微调脚本

微调使用 torchrun + DeepSpeed 进行分布式训练

%%writefile ./src/ds-train-dist.sh
#!/bin/bash
CURRENT_HOST="${SM_CURRENT_HOST}"


IFS=',' read -ra hosts_array <<< "${SM_HOSTS}"
NNODES=${#hosts_array[@]}
NODE_RANK=0

for i in "${!hosts_array[@]}"; do
    if [[ "${hosts_array[$i]}" == *${CURRENT_HOST}* ]]; then
        echo "host index：$i"
        NODE_RANK="$i" 
    fi
done
   
    
MASTER_PORT="13579"
export NCCL_SOCKET_IFNAME="eth0"

#Configure the distributed arguments for torch.distributed.launch.
GPUS_PER_NODE="$SM_NUM_GPUS"
DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE \
                  --nnodes $NNODES \
                  --node_rank $NODE_RANK \
                  --master_addr $MASTER_ADDR \
                  --master_port $MASTER_PORT"

chmod +x ./s5cmd
./s5cmd sync s3://$MODEL_S3_BUCKET/llm/models/baichuan2/baichuan-inc/Baichuan2-13B-Chat/* /tmp/baichuan2/

CMD="torchrun ${DISTRIBUTED_ARGS} ${DEEPSPEED_OPTS}"
echo ${CMD}
${CMD} 2>&1 

if [[ "${CURRENT_HOST}" == "${MASTER_ADDR}" ]]; then  
    ./s5cmd sync /tmp/baichuan2_out s3://$MODEL_S3_BUCKET/llm/models/baichuan2/output/baichuan-inc/Baichuan2-13B-Chat/$(date +%Y-%m-%d-%H-%M-%S)/
fi

启动微调

全参数微调以 Baichuan-2-13B Chat 为例，需要使用至少一台 p4de.12xlarge（8 卡 A100 40GB）作为训练机器

当微调完成后，训练好的模型自动存储于指定的 S3 桶内，可用于后续的模型部署推理

import time
from sagemaker.estimator import Estimator


environment = {
    'MODEL_S3_BUCKET': sagemaker_default_bucket # The bucket to store pretrained model and fine-tune model
}

base_job_name = 'baichuan2-13b-chat-finetune'

instance_type = 'ml.p4d.24xlarge'

estimator = Estimator(role=role,
                      entry_point='ds-train-dist.sh',
                      source_dir='./src',
                      base_job_name=base_job_name,
                      instance_count=1,
                      instance_type=instance_type,
                      image_uri=image_uri,
                      environment=environment,
                      disable_profiler=True,
                      debugger_hook_config=False)


estimator.fit()

总结

大语言模型方兴未艾，正在以各种方式改变和影响着整个世界。客户拥抱大语言模型，亚马逊云科技团队同样在深耕客户需求和大语言模型技术，可以在未来更好的协助客户实现需求、提升业务价值。

亚马逊AWS官方博客