menu Chancel's blog
rss_feed
Chancel's blog
有善始者实繁,能克终者盖寡。

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

作者:Chancel Yang, 创建:2019-03-01, 字数:4560, 已阅:11, 最后更新:2024-06-14

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

Bash
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

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

Bash
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环境

Bash
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套件

Bash
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镜像:

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

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

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

Bash
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. 训练

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

Bash
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

最后,运行训练

Bash
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' 的错误,你可以尝试禁用构建隔离。

Bash
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_utilsfully_sharded_data_parallel文件

Bash
find / -name '_state_dict_utils.py'

修改/path/_state_dict_utils.py

Python
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. 尾声

参考资料


[[replyMessage== null?"发表评论":"发表评论 @ " + replyMessage.m_author]]

account_circle
email
web_asset
textsms

评论列表([[messageResponse.total]])

还没有可以显示的留言...
gravatar
[[messageItem.m_author]] [[messageItem.m_author]]
[[messageItem.create_time]]
[[getEnviron(messageItem.m_environ)]]
[[subMessage.m_author]] [[subMessage.m_author]] @ [[subMessage.parent_message.m_author]] [[subMessage.parent_message.m_author]]
[[subMessage.create_time]]
[[getEnviron(messageItem.m_environ)]]