使用A100GPU训练Vicuna-13b模型的实践

1. 背景

最近我有幸拿得了A100 GPU资源，用于训练AI模型。我进行了一次关于使用A100训练Vicuna-13b模型的实践。

以下是这次训练的环境配置：

8个A100 GPU（每个GPU内存40Gb） 640Gb内存 256核CPU

在进行训练时，可以选择直接在物理机上进行，也可以使用Docker容器。而我选择了使用Docker作为训练工具，下面是软件环境

系统：Ubuntu2204 CUDA：12.2 Docker：20.10.21-0ubuntu1~22.04.3

2. 实践

2.1. 准备

首先，服务器需要安装CUDA。我选择直接安装最新版本的CUDA（12.0），这样可以方便后续在Docker容器中选择不同的低版本CUDA

content_copy

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda-repo-ubuntu2204-12-2-local_12.2.0-535.54.03-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-12-2-local_12.2.0-535.54.03-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-12-2-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

接下来，我们需要克隆仓库并下载模型，这两个步骤需要良好的网络连接支持

content_copy

suao apt install git-lfs

git lfs install

git clone https://huggingface.co/lmsys/vicuna-13b-v1.3
git clone https://github.com/lm-sys/FastChat.git

然后，安装Docker环境

content_copy

sudo apt update
sudo apt upgrade
sudo apt install docker.io
sudo systemctl start docker
sudo systemctl enable docker

# 验证安装是否成功，如成功则输出hello,world
sudo docker run hello-world

安装docker-nvidia套件

content_copy

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-docker.gpg
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list > /dev/null
sudo apt update
sudo apt install nvidia-docker2
sudo systemctl restart docker

# 验证是否安装成功，如成功则输出显卡信息
sudo docker run --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

至此，我们已经准备好了基本所需的工具

2.2. 容器

在可用的镜像中，有许多选择，包括以CUDA为主的镜像和以PyTorch为主的镜像。在这里，我们选择后者，因为前者需要自己编译安装Python和安装PyTorch，会更加复杂一些。我们根据具体的训练场景来进行选择

首先，我们需要拉取CUDA 11.7版本的PyTorch镜像：

content_copy

docker pull pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel

如果需要其他CUDA版本的镜像，请参考：https://hub.docker.com/r/pytorch/pytorch/tags

接下来，我们需要运行容器。根据实际情况填入Fastchat仓库以及模型位置，这里假设它们都放在/data目录下：

content_copy

docker run -itd --gpus all --privileged --cap-add=IPC_LOCK --ulimit memlock=-1 --ulimit stack=67108864 -v /data:/app/ --net=host --ipc=host --name=train-13b pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel

2.3. 训练

进入容器，并切换到仓库目录下，并安装相关依赖：

content_copy

apt update && apt install -y git
cd /app/FastChat
pip install --upgrade pip setuptools wheel -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install einops deepspeed peft -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install flash-attn -i https://pypi.tuna.tsinghua.edu.cn/simple

最后，运行训练

content_copy

torchrun --nproc_per_node=8 --master_port=20001 fastchat/train/train_mem.py \
    --model_name_or_path /app/vicuna-13b-v1.3  \
    --data_path fastchat/data/dummy_conversation.json \
    --output_dir /app/models/vicuna-13b-v1.3-output \
    --bf16 True \
    --num_train_epochs 1 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --evaluation_strategy "steps" \
    --eval_steps 1500 \
    --save_strategy "steps" \
    --save_steps 1500 \
    --save_total_limit 8 \
    --learning_rate 2e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.04 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --fsdp "full_shard auto_wrap offload" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 True \
    --model_max_length 2048 \
    --gradient_checkpointing True \
    --lazy_preprocess True

3. 一些小问题

3.1. No module named 'torch

如果在安装 flash-attn 时遇到 ModuleNotFoundError: No module named 'torch' 的错误，你可以尝试禁用构建隔离。

content_copy

pip install flash-attn --no-build-isolation -i https://pypi.doubanio.com/simple

3.2. CUDA out of memory

如果你在使用 LlamaDecoderLayer 进行训练，并且配置足够保证训练，但仍然遇到 Error: CUDA out of memory.错误，可以考虑以下修复方法

查找_state_dict_utils 或 fully_sharded_data_parallel文件

content_copy

find / -name '_state_dict_utils.py'

修改/path/_state_dict_utils.py

content_copy

state_dict[fqn] = state_dict[fqn].clone().detach()
# 改为
state_dict[fqn] = state_dict[fqn].cpu().clone().detach()

路径参考

torch/distributed/fsdp/_state_dict_utils torch/distributed/fsdp/fully_sharded_data_parallel

4. 尾声

参考资料

discuss.pytorch.org/t/fsdp-failed-to-save-model-checkpoints/178232