LLaMA-Factory 分布式训练实践

一、前置环境准备

1. 基础依赖安装

bash 复制代码

# 更新系统依赖
sudo apt update && sudo apt install -y python3-pip python3-dev git gcc g++ make

# 安装 CUDA（推荐 11.8/12.1，需匹配显卡驱动）
# 方式1：NVIDIA 官方源（以 CUDA 11.8 为例）
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update && sudo apt install -y cuda-11-8

# 配置 CUDA 环境变量（写入 ~/.bashrc）
echo "export PATH=/usr/local/cuda-11.8/bin:\$PATH" >> ~/.bashrc
echo "export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:\$LD_LIBRARY_PATH" >> ~/.bashrc
source ~/.bashrc

# 验证 CUDA
nvcc -V  # 应输出 CUDA 版本信息
nvidia-smi  # 应显示显卡列表

2. LLaMA-Factory 安装

bash 复制代码

# 克隆仓库（建议用最新版）
git clone https://github.com/hiyouga/LLaMA-Factory.git
cd LLaMA-Factory

# 创建虚拟环境（可选但推荐）
conda create -n llama_factory python=3.10 -y
conda activate llama_factory

# 安装依赖（包含分布式训练所需依赖）
pip install -e .[deepspeed,torch]  # deepspeed 依赖需单独指定
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118  # 匹配 CUDA 版本

# 验证安装
llama_factory-cli version

3. 分布式训练核心依赖检查

bash 复制代码

# 检查 torch.distributed 可用性
python -c "import torch; print(torch.distributed.is_available())"  # 输出 True
# 检查 DeepSpeed
deepspeed --version  # 输出 DeepSpeed 版本
# 检查 FSDP（PyTorch 内置，需 PyTorch >= 2.0）
python -c "import torch.distributed.fsdp; print('FSDP available')"

二、核心概念说明

分布式引擎	核心特点	适用场景
DDP（Distributed Data Parallel）	PyTorch 原生，简单稳定，显存占用略高	单机多卡/小规模多机，模型 < 10B
DeepSpeed	支持 ZeRO 优化（1/2/3 阶段），显存效率高，功能丰富	大模型（10B+）、显存受限场景
FSDP（Fully Sharded Data Parallel）	PyTorch 2.0+ 主推，原生支持模型分片，兼容 torch.compile	大模型、多机多卡，追求原生兼容性

三、单机多卡训练

1. DDP 单机多卡

核心配置

无需额外配置文件，通过 torchrun 启动即可，LLaMA-Factory 会自动适配 DDP。

启动命令

bash 复制代码

# 示例：4卡训练 Llama-2-7B，使用 DDP
torchrun --nproc_per_node=4 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /path/to/llama-2-7b \
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir ./output/llama-2-7b-ddp \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --fp16 \  # 混合精度，降低显存占用
  --ddp_find_unused_parameters False \  # DDP 优化
  --plot_loss \
  --disable_tqdm False

2. DeepSpeed 单机多卡

步骤1：编写 DeepSpeed 配置文件

创建 ds_config.json（示例为 ZeRO-2 配置，显存优化效果好）：

json 复制代码

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4,
  "bf16": {
    "enabled": true  # 支持 BF16 的显卡（如 A100/V100）建议开启，否则用 fp16
  },
  "fp16": {
    "enabled": false
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 5e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupCosineLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 5e-5,
      "warmup_num_steps": 100,
      "total_num_steps": 10000
    }
  },
  "gradient_clipping": 1.0,
  "logging": {
    "steps_per_print": 10
  },
  "offload_param": {
    "device": "cpu",
    "pin_memory": true
  },
  "zero_allow_untested_optimizer": true,
  "communication_data_type": "float16",
  "enable_stage3_param_sharding": true
}

步骤2：启动训练

bash 复制代码

deepspeed --num_gpus=4 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /path/to/llama-2-7b \
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir ./output/llama-2-7b-deepspeed \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --deepspeed ./ds_config.json \  # 指定 DeepSpeed 配置文件
  --plot_loss \
  --disable_tqdm False

3. FSDP 单机多卡

核心配置：通过参数指定 FSDP 选项

bash 复制代码

torchrun --nproc_per_node=4 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /path/to/llama-2-7b \
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir ./output/llama-2-7b-fsdp \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --fp16 \
  --fsdp "full_shard auto_wrap" \  # FSDP 核心配置：全分片 + 自动包装模型
  --fsdp_transformer_layer_cls_to_wrap "LlamaDecoderLayer" \  # 适配 Llama 模型
  --fsdp_use_orig_params True \  # 兼容 LoRA
  --plot_loss \
  --disable_tqdm False

四、多机多卡训练

前提条件

所有机器处于同一局域网，可互相 ping 通；
所有机器的 CUDA、PyTorch、LLaMA-Factory 版本一致；
模型、数据集路径在所有机器上保持一致（建议用 NFS 共享存储）；
配置免密 SSH 登录（主节点可免密登录从节点）。

1. DDP 多机多卡

步骤1：配置主机列表

创建 hosts.txt（示例：2 机，每机 4 卡）：

复制代码

192.168.1.100 slots=4  # 主节点 IP + 卡数
192.168.1.101 slots=4  # 从节点 IP + 卡数

步骤2：启动训练（主节点执行）

bash 复制代码

torchrun --nnodes=2 \
  --node_rank=0 \  # 主节点 rank 为 0
  --master_addr=192.168.1.100 \  # 主节点 IP
  --master_port=29500 \  # 主节点端口（未被占用）
  --nproc_per_node=4 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /nfs/llama-2-7b \  # NFS 共享路径
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir /nfs/output/llama-2-7b-ddp-multi \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --fp16 \
  --ddp_find_unused_parameters False \
  --plot_loss

从节点启动（可选，若未自动连接）

bash 复制代码

torchrun --nnodes=2 \
  --node_rank=1 \
  --master_addr=192.168.1.100 \
  --master_port=29500 \
  --nproc_per_node=4 \
  src/train_bash.py \
  # 其余参数与主节点完全一致

2. DeepSpeed 多机多卡

步骤1：修改 DeepSpeed 配置文件

ds_config.json 无需大幅修改，仅需确保 train_batch_size = 总卡数 × per_gpu_batch × accumulation_steps。

步骤2：启动训练（主节点执行）

bash 复制代码

deepspeed --num_nodes=2 \
  --num_gpus=4 \
  --master_addr=192.168.1.100 \
  --master_port=29500 \
  --node_rank=0 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /nfs/llama-2-7b \
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir /nfs/output/llama-2-7b-deepspeed-multi \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --deepspeed ./ds_config.json \
  --plot_loss

3. FSDP 多机多卡

与 DDP 多机启动方式一致，仅需保留 FSDP 相关参数：

bash 复制代码

torchrun --nnodes=2 \
  --node_rank=0 \
  --master_addr=192.168.1.100 \
  --master_port=29500 \
  --nproc_per_node=4 \
  src/train_bash.py \
  --stage sft \
  --do_train \
  --model_name_or_path /nfs/llama-2-7b \
  --dataset alpaca_gpt4_en \
  --finetuning_type lora \
  --output_dir /nfs/output/llama-2-7b-fsdp-multi \
  --overwrite_cache \
  --per_device_train_batch_size 4 \
  --gradient_accumulation_steps 4 \
  --lr_scheduler_type cosine \
  --logging_steps 10 \
  --save_steps 1000 \
  --learning_rate 5e-5 \
  --num_train_epochs 3.0 \
  --fp16 \
  --fsdp "full_shard auto_wrap" \
  --fsdp_transformer_layer_cls_to_wrap "LlamaDecoderLayer" \
  --fsdp_use_orig_params True \
  --plot_loss

五、常见问题解决

1. 显存不足

降低 per_device_train_batch_size，增大 gradient_accumulation_steps；
开启 fp16/bf16 混合精度；
DeepSpeed 升级 ZeRO 阶段（如 ZeRO-3），或开启参数/优化器 offload 到 CPU；
FSDP 增加 --fsdp_offload_params_to_cpu True。

2. 多机通信失败

检查防火墙/端口是否开放（master_port 需放行）；
确保所有机器的 PyTorch 分布式端口未被占用；
验证 SSH 免密登录是否生效；
主节点 master_addr 需用内网 IP，而非 127.0.0.1。

3. DeepSpeed 报错 "ZeRO stage 3 not supported"

确保 DeepSpeed 版本 ≥ 0.9.0；
配置文件中添加 "zero_allow_untested_optimizer": true；
避免 LoRA 与 ZeRO-3 同时使用（部分场景不兼容）。

4. FSDP 与 LoRA 不兼容

必须添加 --fsdp_use_orig_params True；
PyTorch 版本需 ≥ 2.0.0；
确保 fsdp_transformer_layer_cls_to_wrap 匹配模型（如 ChatGLM 为 GLMBlock）。

六、性能对比（参考）

引擎	单机4卡（A100）	多机8卡（A100）	显存占用（7B模型）	训练速度（tokens/s）
DDP	稳定	较稳定	~18GB/卡	~800
DeepSpeed（ZeRO-2）	稳定	稳定	~12GB/卡	~750（略慢但显存省）
FSDP	稳定（PyTorch2.0+）	较稳定	~10GB/卡	~780

总结

快速上手：优先选择 DDP，配置简单、稳定性高；
大模型训练：优先选择 DeepSpeed（ZeRO-2/3）或 FSDP，显存效率更高；
多机训练：DDP/FSDP 配置更简洁，DeepSpeed 需确保多机通信正常；
所有场景建议使用 LoRA 微调（而非全量微调），大幅降低显存占用。