最近我有幸拿得了A100 GPU资源，用于训练AI模型。我进行了一次关于使用A100训练Vicuna-13b模型的实践。
以下是这次训练的环境配置：
在进行训练时，可以选择直接在物理机上进行，也可以使用Docker容器。而我选择了使用Docker作为训练工具，下面是软件环境
首先，服务器需要安装CUDA。我选择直接安装最新版本的CUDA（12.0），这样可以方便后续在Docker容器中选择不同的低版本CUDA
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
接下来，我们需要克隆仓库并下载模型，这两个步骤需要良好的网络连接支持
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环境
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套件
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
至此，我们已经准备好了基本所需的工具
在可用的镜像中，有许多选择，包括以CUDA为主的镜像和以PyTorch为主的镜像。在这里，我们选择后者，因为前者需要自己编译安装Python和安装PyTorch，会更加复杂一些。我们根据具体的训练场景来进行选择
首先，我们需要拉取CUDA 11.7版本的PyTorch镜像：
docker pull pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
如果需要其他CUDA版本的镜像，请参考：https://hub.docker.com/r/pytorch/pytorch/tags
接下来，我们需要运行容器。根据实际情况填入Fastchat仓库以及模型位置，这里假设它们都放在/data目录下：
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
进入容器，并切换到仓库目录下，并安装相关依赖：
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
最后，运行训练
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
如果在安装 flash-attn 时遇到 ModuleNotFoundError: No module named 'torch' 的错误，你可以尝试禁用构建隔离。
pip install flash-attn --no-build-isolation -i https://pypi.doubanio.com/simple
如果你在使用 LlamaDecoderLayer 进行训练，并且配置足够保证训练，但仍然遇到 Error: CUDA out of memory.错误，可以考虑以下修复方法
查找_state_dict_utils 或 fully_sharded_data_parallel文件
find / -name '_state_dict_utils.py'
修改/path/_state_dict_utils.py
state_dict[fqn] = state_dict[fqn].clone().detach()
# 改为
state_dict[fqn] = state_dict[fqn].cpu().clone().detach()
路径参考
torch/distributed/fsdp/_state_dict_utilstorch/distributed/fsdp/fully_sharded_data_parallel参考资料