使用 LLaMA-Factory 进行训练

LLaMA-Factory 是一个统一高效的大语言模型微调框架，支持多种微调方法（LoRA、QLoRA、全参数微调等）和分布式训练。本文详细描述了如何使用 AIStudio 「任务」功能配置和启动 LLaMA-Factory 分布式训练任务。

准备自定义训练镜像

由于 AIStudio 任务是随启随用的方式，不会长期占用服务器，所以我们每次启动一个训练任务的时候，AIStudio 控制平面会去申请服务器，然后安装训练代码并运行，为了提高训练准备的效率，我们需要把 LLaMA Factory 以及训练代码提前构建成为一个训练镜像，并且推送到镜像中心，然后在训练任务创建时指定这个镜像地址，AIStudio 启动任务时会下载这个镜像到服务器中，通过这种方式我们可以使用任何算法或库，极大的提高了的灵活性。

您可以通过以下方式使用自定义训练镜像。

• 获取 LLaMA Factory 官方镜像，存入镜像中心

  # 该步骤要求需要配置代理或使用第三方加速服务
  docker pull --platform linux/amd64 hiyouga/llamafactory:0.9.4

  具体版本详见 LLaMA Factory Docker Hub 镜像列表。关于如何存入镜像中心，详见迁移外部镜像至租户镜像仓库。

• 自行构建 LLaMA Factory 镜像。在 AIStudio 镜像中心可使用 Dockerfile 构建镜像。详见构建自定义镜像。

  Dockerfile 示例：

  FROM cr.infini-ai.com/infini-ai/ubuntu:22.04-20240429
  
  # 设置环境变量
  # DEBIAN_FRONTEND=noninteractive 防止安装过程中的交互式提示
  # PATH 添加CUDA可执行文件路径
  # LD_LIBRARY_PATH 添加CUDA库文件路径
  # CUDA_HOME 设置CUDA根目录
  # USE_MODELSCOPE_HUB 启用ModelScope模型仓库
  ENV DEBIAN_FRONTEND=noninteractive \
      PATH=/usr/local/cuda/bin:/usr/local/cuda-12.2/bin${PATH:+:${PATH}} \
      LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.2/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} \
      CUDA_HOME=/usr/local/cuda \
      USE_MODELSCOPE_HUB=1
  
  # 安装CUDA并配置系统环境
  # 1. 更新apt并安装必要工具
  # 2. 清理apt缓存减小镜像体积
  # 3. 下载CUDA安装包
  # 4. 静默安装CUDA工具包（请勿安装 driver）
  # 5. 配置动态链接库搜索路径
  # 6. 更新动态链接库缓存
  # 7. 清理安装包
  RUN apt-get update && apt-get install -y wget build-essential \
      && rm -rf /var/lib/apt/lists/* \
      && wget https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run \
      && sh cuda_12.2.2_535.104.05_linux.run --toolkit --silent --override \
      && echo '/usr/local/cuda/lib64' >> /etc/ld.so.conf \
      && echo '/usr/local/cuda-12.2/lib64' >> /etc/ld.so.conf \
      && ldconfig \
      && rm cuda_12.2.2_535.104.05_linux.run
  
  # 安装PyTorch/LLaMA-Factory/ModelScope/Tensorboard 等
  # 1. 安装指定版本的PyTorch
  # 2. 安装 Tensorboard
  # 3. 克隆LLaMA-Factory代码库（使用--depth 1只克隆最新版本以减小体积）
  # 4. 切换到LLaMA-Factory目录
  # 5. 以可编辑模式安装LLaMA-Factory，并包含torch/metrics/modelscope附加功能
  RUN pip install --no-cache-dir torch==2.4.0 \
      && pip install --no-cache-dir tensorboard \
      && git clone --depth 1 https://gh-proxy.com/https://github.com/hiyouga/LLaMA-Factory.git \
      && cd LLaMA-Factory \
      && pip install --no-cache-dir -e ".[torch,metrics,modelscope]"

关于如何获取 LLaMA Factory 镜像，可参考以下教程：

📖 快捷安装 LLaMA Factory

基本流程

在 AIStudio 任务中，通过以下步骤启动 LLaMA-Factory 分布式训练任务：

1. 创建训练任务：配置任务 Worker 规格、Worker 数量，启动命令、分布式框架等。
2. 初始化：平台创建对应的 pod，并注入必要环境变量，供用户代码获取环境信息。
3. 容错检查：训练开始前，平台发起自检，确保所有 Worker 之间网络通畅、GPU、存储工作正常。
4. 训练执行：
   • 若任意 pod 失败（退出码非 0），训练任务失败。
   • 若所有 pod 成功完成（退出码为 0），训练任务成功。

关键配置

• Huggingface 官方地址从国内访问不太稳定，因此我们需要通过 export HF_ENDPOINT=https://hf-mirror.com 设置镜像站，提供模型和数据集下载加速能力。
• Huggingface Hub 或 ModelScope 默认将所下载的模型和数据集放在主目录下（例如 ~/.cache/huggingface），但这个是本地目录，无法让分布式训练中多个 worker 共享，因此我们需要修改该地址为共享文件存储的位置。
  • Huggingface Hub 需要配置 HF_HOME
  • ModelScope 需要配置 MODELSCOPE_CACHE

启动命令

「启动命令」是容器启动时执行的入口脚本，用于配置和启动模型训练任务。您需要在此字段中提供完整的 Shell 命令或脚本，用于：

• 配置训练环境：设置环境变量、安装依赖包等
• 启动分布式训练：使用 torchrun、mpirun 等命令启动训练

单机多卡训练

在创建任务时，请在分布式框架处指定为单机。

以下命令用于启动单机多卡的 LLaMA-Factory 训练，可以直接复制到「启动命令」输入框：

#!/bin/bash

set -o pipefail

### 官方 LLaMA Factory 镜像 配置使用清华 pip 源 ###
pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
### 官方 LLaMA Factory 镜像 安装 Tensorboard ###
pip install tensorboard

################## 分布式 DDP 训练配置 ##################
GPUS_PER_NODE=${GPUS_PER_NODE:-2}
# Single-node setup
NNODES=1
NODE_RANK=0

DISTRIBUTED_ARGS=(
    --standalone
    --nnodes $NNODES 
    --nproc_per_node $GPUS_PER_NODE
)

################## Huggingface 配置 ##################
# 配置 hf 国内地址
export HF_ENDPOINT=${HF_ENDPOINT:-"https://hf-mirror.com"}
# 配置 Huggingface home 到共享存储，使得多个 worker 可以共同访问
export HF_HOME=${HF_HOME:-"/mnt/zhaoyinghao/huggingface"}

################## 模型配置 ##################
TEMPLATE="llama3"
MODEL_NAME_OR_PATH="/mnt/zhaoyinghao/Meta-Llama-3-8B-Instruct/"
DATASET="alpaca_en_demo"

################## 输出配置 ##################
OUTPUT_DIR=${OUTPUT_DIR:-"/mnt/zhaoyinghao/saves/llama3-8b/lora/"}
mkdir -p ${OUTPUT_DIR}

# 如果使用 LLaMA Factory 官方镜像，进入 /app 目录
cd /app

################## 训练脚本配置 ##################
# 使用 tee 命令使日志写入到标准输出的同时写入到日志文件中

torchrun ${DISTRIBUTED_ARGS[@]} \
    src/train.py \
    --stage sft \
    --model_name_or_path ${MODEL_NAME_OR_PATH}  \
    --do_train \
    --dataset ${DATASET} \
    --template ${TEMPLATE} \
    --finetuning_type lora \
    --output_dir ${OUTPUT_DIR} \
    --overwrite_cache \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --lr_scheduler_type cosine \
    --logging_steps 100 \
    --save_steps 500 \
    --learning_rate 1e-4 \
    --num_train_epochs 20.0 \
    --plot_loss \
    --report_to tensorboard \
    --bf16 2>&1 | tee ${OUTPUT_DIR}/node-${NODE_RANK}.log

# 获取上一条命令的返回值，便于排查问题
ret=$?
if [[ ${ret} -ne 0 ]]; then
    echo "[$(date +"%Y-%m-%d %H:%M:%S")] $HOSTNAME 训练失败！现场保留时间为 10000s"
    sleep 10000
fi

exit $ret

多机多卡训练

在创建任务时，请在分布式框架处指定为 Pytorch DDP。

注意

• 建议使用支持 IB/RoCE 训练网的资源规格。如果无训练网，建议在启动命令中配置 export NCCL_SOCKET_IFNAME=eth0。
• 如使用 RoCE 类型训练网规格，建议在规格包含卡数低于 8 卡关闭 RDMA 配置。

对于多机多卡训练，需要使用平台预置的环境变量来配置分布式参数：

#!/bin/bash

set -o pipefail

### 官方 LLaMA Factory 镜像 配置使用清华 pip 源 ###
pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
### 官方 LLaMA Factory 镜像 安装 Tensorboard ###
pip install tensorboard

################## 分布式 DDP 训练配置 ##################
# 获取每个 Worker 的 GPU 数量
GPUS_PER_NODE=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)

# 设置分布式训练参数
MASTER_ADDR=${MASTER_ADDR:-'127.0.0.1'}
MASTER_PORT=${MASTER_PORT:-'29500'}
NNODES=${WORLD_SIZE:-'1'}  # 平台提供的 WORLD_SIZE 是节点数
NODE_RANK=${RANK:-'0'}     # 平台提供的 RANK 是 pod 编号，用于 --node_rank

DISTRIBUTED_ARGS=(
    --nproc_per_node $GPUS_PER_NODE
    --nnodes $NNODES
    --node_rank $NODE_RANK
    --master_addr $MASTER_ADDR
    --master_port $MASTER_PORT
)

################## Huggingface 配置 ##################
# 配置 hf 国内地址
export HF_ENDPOINT=${HF_ENDPOINT:-"https://hf-mirror.com"}
# 配置 Huggingface home 到共享存储，使得多个 worker 可以共同访问
export HF_HOME=${HF_HOME:-"/mnt/zhaoyinghao/huggingface"}

################## 模型配置 ##################
TEMPLATE="llama3"
MODEL_NAME_OR_PATH="/mnt/zhaoyinghao/Meta-Llama-3-8B-Instruct/"
DATASET="alpaca_en_demo"

################## 输出配置 ##################
OUTPUT_DIR=${OUTPUT_DIR:-"/mnt/zhaoyinghao/saves/llama3-8b/lora/"}
mkdir -p ${OUTPUT_DIR}

################## 训练脚本配置 ##################
# 使用 tee 命令使日志写入到标准输出的同时写入到日志文件中

cd /app

torchrun ${DISTRIBUTED_ARGS[@]} \
    src/train.py \
    --stage sft \
    --model_name_or_path ${MODEL_NAME_OR_PATH}  \
    --do_train \
    --dataset ${DATASET} \
    --template ${TEMPLATE} \
    --finetuning_type lora \
    --output_dir ${OUTPUT_DIR} \
    --overwrite_cache \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --lr_scheduler_type cosine \
    --logging_steps 100 \
    --save_steps 500 \
    --learning_rate 1e-4 \
    --num_train_epochs 20.0 \
    --plot_loss \
    --report_to tensorboard \
    --bf16 2>&1 | tee ${OUTPUT_DIR}/node-${NODE_RANK}.log

# 获取上一条命令的返回值，便于排查问题
ret=$?
if [[ ${ret} -ne 0 ]]; then
    echo "[$(date +"%Y-%m-%d %H:%M:%S")] $HOSTNAME 训练失败！现场保留时间为 10000s"
    sleep 10000
fi

exit $ret

注意

您可以根据实际需求修改脚本中的配置参数，如模型路径、数据集、输出目录等。建议使用动态检测以提高灵活性。

理解平台内置环境变量

AIStudio 平台在分布式训练任务中采用了与 PyTorch 原生定义不同的环境变量设置。以下是平台提供的内置环境变量：

• MASTER_ADDR：DDP 分布式通信的主机 IP 地址或名称（通常为 worker0 的 pod 名称），由系统自动解析。
• MASTER_PORT：主节点上开启的端口号，默认值为 29500。
• WORLD_SIZE：启动的 pod 数量（即节点数量），注意：不同于 PyTorch 原生的 WORLD_SIZE（总进程数）。
• RANK：每个 pod 的编号（从 0 到 WORLD_SIZE - 1），用于设置 torchrun 的 --node_rank 参数。

注意

平台将 WORLD_SIZE 定义为 pod（节点）的数量，而非 PyTorch 中的总进程数。这一选择与 Kubeflow 的 pytorch-operator 保持一致，方便从 Kubeflow 迁移的用户无缝过渡。

LLaMA-Factory DDP 关键参数

LLaMA-Factory 通过 torchrun 启动分布式训练，需要配置以下参数：

• --master_addr：主节点地址，对应平台注入的环境变量 MASTER_ADDR。

• --master_port：主节点端口，对应平台注入的环境变量 MASTER_PORT。

• --nnodes：节点数量，对应平台注入的环境变量 WORLD_SIZE。

• --node_rank：节点编号，对应平台注入的环境变量 RANK。

• --nproc_per_node：每个 pod 的进程数，通常等于 GPU 数量。例如：

  GPUS_PER_NODE=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)

LLaMA-Factory 的主要训练参数通过命令行传递给 src/train.py：

• --stage：训练阶段（sft、dpo、ppo、rm 等）
• --model_name_or_path：预训练模型路径
• --finetuning_type：微调类型（lora、full 等）
• --dataset：训练数据集名称
• --template：对话模板（如 llama3）
• --output_dir：输出目录

故障排除

1. 分布式训练未启动

   • 检查 torchrun 参数配置是否正确
   • 环境变量 MASTER_ADDR、MASTER_PORT 等设置

2. 内存不足错误

   • 减小 --per_device_train_batch_size
   • 增加 --gradient_accumulation_steps
   • 考虑使用 LoRA 而非全参数微调

3. 模型加载失败

   • 检查 --model_name_or_path 路径是否正确
   • 确保共享存储挂载正常
   • 验证 HuggingFace 配置

4. 数据集加载失败

   • 确认 --dataset 名称是否正确
   • 检查数据集是否在 LLaMA-Factory 支持列表中

调试工具

• 在任务运行过程中，可从网页端登录任务 Worker，访问 Web Terminal。在任务详情页底部可看到登录按钮。
• 任务功能提供 atlctl 命令行调试工具，您可以从 Web Terminal 登录任意的任务 Worker，执行停止任务、统一下发测试命令等调试工作。