使用ModelScope Swift微调Baichuan2模型

介绍

Baichuan 2是百川智能推出的开源大语言模型，采用2.6万亿Tokens的高质量语料进行训练，在多个权威的中文、英文和多语言的通用、领域benchmark上取得了同尺寸最佳的效果。Baichuan2 目前发布了7B、13B的Base和Chat版本，支持模型商用。

当在特定领域使用大语言模型时，可以通过prompt的方式引导模型，也可以通过在领域数据集上微调训练，从而在领域的任务上获得更好的效果。后者的优点是不依赖于Prompt（可能超过模型的输入长度上限），有更好的推理性能，并且经过微调后，在领域相关任务上有更好的效果。

本文将介绍如何在PAI对Baichuan2模型完成微调训练。

安装和配置SDK

我们需要首先安装PAI Python SDK以运行本示例。

!python -m pip install --upgrade alipai

SDK需要配置访问阿里云服务需要的AccessKey，以及当前使用的工作空间和OSS Bucket。在PAI SDK安装之后，通过在命令行终端 中执行以下命令，按照引导配置密钥、工作空间等信息。

# 以下命令，请在 命令行终端 中执行.

python -m pai.toolkit.config

我们可以通过以下代码验证配置是否已生效。

import pai
from pai.session import get_default_session

print(pai.__version__)

sess = get_default_session()

# 获取配置的工作空间信息
assert sess.workspace_name is not None
print(sess.workspace_name)

准备训练脚本

ModelScope提供了SWIFT(Scalable lightWeight Infrastructure for Fine-Tuning)框架，支持模型的全参数微调，也集成了各种高效微调方法，例如LoRA、QLoRA等，支持用户对Baichuan2、QWen、llama2等常见的语言进行微调训练。

基于Swift的LLM finetune脚本，我们修改了部分逻辑，从而支持用户在PAI的训练作业中使用，主要包括：

• 使用PAI预置的Baichuan2-Base模型

对于热门的社区模型，PAI提供了模型缓存在OSS Bucket上，支持挂载到训练作业，训练脚本可以通过读取本地文件的方式加载获取模型。

• 保存模型

训练脚本需要将模型保存到指定路径(/ml/output/model)，从而将模型保存到用户的OSS Bucket中。

• 训练依赖的第三方

训练作业将运行在PAI提供的PyTorch基础镜像上，我们需要在作业环境中安装transformers、datasets、swift、xformers等第三方依赖。PAI训练作业支持使用训练脚本目录下的requirements.txt安装第三方依赖。

完整的训练脚本请参考 train_src 目录下的训练文件(llm_sft.py)。

提交训练作业

使用提交任务的方式训练模型，能够支持用户并行运行多个训练任务，高效得探索不同的超参组合对于模型性能影响，并且能够支持分布式训练。通过PAI Python SDK提供的EstimatorAPI，我们可以方便得将一个本地训练脚本提交到PAI上运行。

我们将通过以下代码配置训练作业脚本、作业启动命令、使用的作业镜像，以及机器实例规格，提交训练作业。

from pai.image import retrieve
from pai.estimator import Estimator

# 训练作业启动命令
# 完整的参数说明请参考文档：https://github.com/modelscope/swift/blob/main/examples/pytorch/llm/README_CN.md#sftsh-%E5%91%BD%E4%BB%A4%E8%A1%8C%E5%8F%82%E6%95%B0
command = """CUDA_VISIBLE_DEVICES=0 \
python llm_sft.py \
    --model_type baichuan2-7b \
    --sft_type lora \
    --template_type default-generation \
    --dtype fp16 \
    --output_dir /ml/output/model/ \
    --dataset advertise-gen \
    --train_dataset_sample 20000 \
    --num_train_epochs 1 \
    --max_length 2048 \
    --quantization_bit 4 \
    --lora_rank 8 \
    --lora_alpha 32 \
    --lora_dropout_p 0. \
    --lora_target_modules ALL \
    --gradient_checkpointing true \
    --batch_size 16 \
    --weight_decay 0. \
    --learning_rate 1e-4 \
    --gradient_accumulation_steps 4 \
    --max_grad_norm 0.5 \
    --warmup_ratio 0.03 \
    --eval_steps 100 \
    --save_steps 100 \
    --save_total_limit 2 \
    --logging_steps 10
"""


# 配置训练作业
est = Estimator(
    source_dir="train_src/",  # 代码目录
    image_uri=retrieve("PyTorch", framework_version="latest").image_uri,  # 训练作业使用的镜像
    command=command,  # 训练启动命令
    instance_type="ecs.gn6e-c12g1.3xlarge",  # 使用的机器规格示例，V100(32G)
    instance_count=1,  # 机器实例个数
    base_job_name="baichuan2_finetune",  # 训练作业名称
)

PAI提供了预置的Baichuan2-Base模型，可以通过以下方式获取对应的模型OSS Bucket路径。

from pai.model import RegisteredModel

# 获取PAI提供的Baichuan2-7B-Base模型
m = RegisteredModel(
    model_name="baichuan-inc/Baichuan2-7B-Base", model_provider="huggingface"
)

# 模型地址
print(m.model_data)

提交训练作业，等待作业完成。用户可以通过打印的作业详情页URL，查看训练作业进度，资源使用，日志等信息。

# 提交训练作业
est.fit(
    inputs={
        # 训练代码可以从 /ml/input/data/pretrained_model/ 目录下读取挂载的预训练模型
        "pretrained_model": m.model_data,
    },
    wait=False,  # 是否等待训练作业完成
)

# 打开一个TensorBoard，监控训练作业
tb = est.tensorboard()


# 等待训练作业完成
est.wait()

训练作业写出到 /ml/output/model 目录下的模型文件和checkpoints将被保存到用户的OSS Bucket中，可以通过 est.model_data() 获取 OSS Bucket路径。

# 查看数据模型的OSS Bucket路径
print(est.model_data())


# 删除启动的TensorBoard（每一个账号下最多能够启动5个TensorBoard示例）
tb.delete()

结语

在当前示例中，我们展示了如何基于ModelScope Swift框架，使用PAI预置的Baichuan2-Base模型，完成Baichuan2模型的微调训练。用户可以参考以上的示例，修改脚本，使用用户自定义的数据集，或是修改使用的基础预训练模型，完成自定义语言模型的微调训练。