Docker容器搭建单机多卡分布式训练环境

来吧互相伤害吧

已于 2024-06-12 16:57:18 修改

阅读量1.2k

点赞数 21

文章标签： docker 分布式 pytorch

于 2024-04-30 21:05:12 首次发布

本文链接：https://blog.youkuaiyun.com/m0_50917429/article/details/138353381

版权

本文详细描述了如何解决CUDA版本过高导致的deepspeed脚本启动问题，通过Docker容器配置适当CUDA版本，挂载存储，更换镜像源，安装conda和torch，并实现环境保存与迁移的过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

# 写在前面

宿主机CUDA版本过高导致无法启动 deepspeed 分布式训练脚本，缺失了很多 CUDA 编译的动态链接库，为了更加方便切换 CUDA 版本，记录使用 Docker 容器搭建单机多卡分布式的环境。

0 更换Docker存储目录

由于宿主机根目录空间可能较小，把Docker存储目录放到挂载盘下面，操作如下：

 vim /etc/docker/daemon.json

# 以目标存储路径为/run/docker_root为例，添加如下内容
{
  "data-root": "/mnt/docker_root"
}

# :wq!保存后重启docker服务
systemctl restart docker

重启后使用 docker info 可以看到 data-root 值已经被修改。

1 拉取Docker镜像并启动

nvidia/cuda - Docker Image | Docker Hubhttps://hub.docker.com/r/nvidia/cuda/

在 Docker Hub 上选取合适的镜像，这里选取版本 11.8.0-cudnn8-devel-ubuntu22.04，pull 到宿主机本地，拉取后启动容器，启动命令如下：

# --gpus 指定显卡的数量
# --shm-size 指定分配内存，尽可能多，便于单机多卡之间通信
# -v 映射目录，该目录和宿主机共享，容器中数据时优先存储与此，便于commit时减小镜像的占用空间
docker run -dit --net=host --name=cuda118 --gpus '"device=0,1,2,3,4,5,6,7"'  --restart=always --shm-size 1024G -v /mnt/hpfs/:/mnt/hpfs/ nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 bash

# 进入容器内部
docker exec -it cuda118 bash

网络使用了 host 模式，和宿主机之间共享所有端口，因此需要更改容器内部的 SSH 端口，否则无法启动容器内的 SSH 端口，详见下一节。

Tips:

如果启动时出现：

docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]].

请参考链接解决：

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.htmlhttps://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html

2 更换镜像源并启动SSH

首先 apt-get update 更新一下安装包，随后 apt-get install vim（顺序不能反，否则无法安装）。完成后换成中科大的镜像源：

# 备份原始安装列表文件
mv /etc/apt/sources.list /etc/apt/sources.list.bk
# 打开新的安装列表文件
vim /etc/apt/sources.list

####################### 粘贴以下内容 #######################

deb https://mirrors.ustc.edu.cn/ubuntu/ jammy main restricted universe multiverse
# deb-src https://mirrors.ustc.edu.cn/ubuntu/ jammy main restricted universe multiverse

deb https://mirrors.ustc.edu.cn/ubuntu/ jammy-security main restricted universe multiverse
# deb-src https://mirrors.ustc.edu.cn/ubuntu/ jammy-security main restricted universe multiverse

deb https://mirrors.ustc.edu.cn/ubuntu/ jammy-updates main restricted universe multiverse
# deb-src https://mirrors.ustc.edu.cn/ubuntu/ jammy-updates main restricted universe multiverse

deb https://mirrors.ustc.edu.cn/ubuntu/ jammy-backports main restricted universe multiverse
# deb-src https://mirrors.ustc.edu.cn/ubuntu/ jammy-backports main restricted universe multiverse

# deb https://mirrors.ustc.edu.cn/ubuntu/ jammy-proposed main restricted universe multiverse
# deb-src https://mirrors.ustc.edu.cn/ubuntu/ jammy-proposed main restricted universe multiverse

############################################################

# :wq!保存后更新安装包源
apt-get update
# 安装 openssh-server
apt-get install openssh-server

安装 openssh-server 后需要调整一下ssh服务配置，操作如下：

vim /etc/ssh/sshd_config

# 修改Port为其他端口，以2211为例
Port 2211
# 允许root登录
PermitRootLogin yes
# :wq!保存后，重启SSH服务
service ssh restart

记得使用 ufw allow 2211 开放防火墙端口，这样可以通过宿主机IP + 更换后的端口号直接SSH访问容器内部。

3 安装conda同时更换镜像源

首先安装 miniconda，把 miniconda 安装到了映射目录下，操作如下：

mkdir -p /mnt/hpfs/miniconda3

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /mnt/hpfs/miniconda3/miniconda.sh

bash ~/miniconda3/miniconda.sh -b -u -p /mnt/hpfs/miniconda3

rm -rf /mnt/hpfs/miniconda3/miniconda.sh

/mnt/hpfs/miniconda3/bin/conda init bash

安装完成后重启一个终端，更换 conda 和 pip 的镜像源，操作如下：

# 打开 conda 配置文件
vim ~/.condarc 

####################### 粘贴以下内容 #######################

channels:
  - defaults
show_channel_urls: true
default_channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/pro
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
custom_channels:
  conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud

############################################################

# 设置 pip 镜像源
pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple

Tips:

安装完 miniconda 后 ~/.bashrc 文件中会自动添加下面几行 conda 的自动初始化脚本：

# >>> conda initialize >>>
# !! Contents within this block are managed by 'conda init' !!
__conda_setup="$('/mnt/hpfs/miniconda3/bin/conda' 'shell.bash' 'hook' 2> /dev/null)"
if [ $? -eq 0 ]; then
    eval "$__conda_setup"
else
    if [ -f "/mnt/hpfs/miniconda3/etc/profile.d/conda.sh" ]; then
        . "/mnt/hpfs/miniconda3/etc/profile.d/conda.sh"
    else
        export PATH="/mnt/hpfs/miniconda3/bin:$PATH"
    fi
fi
unset __conda_setup
# <<< conda initialize <<<

把自动初始化脚本注释掉，使用以下声明来代替，同时 source ~/.bashrc 使更改生效。

alias mini='source /mnt/hpfs/miniconda3/bin/activate'

这样在每次进入容器内部终端后手动输入 mini 即可激活 miniconda 的环境，更加灵活。

4 创建虚拟环境安装torch

# 创建 python 版本为3.10名为deepspeed的虚拟环境
conda create -n deepspeed python=3.10

# 激活并安装torch
conda activate deepspeed
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 进入python检验是否安装成功
python
>>> import torch
>>> torch.cuda.is_available()
>>> torch.version.cuda


# 容器内安装 flash_attention
# FLASH_ATTENTION_FORCE_BUILD=TRUE pip install flash-attn

5 训练环境保存与迁移

以正在运行的容器cuda118为例，通过如下命令进行训练环境的保存：

# 保存运行时的容器为镜像，标签为 11.8.0-cudnn8-devel-ubuntu22.04-v1.0
docker commit cuda118 nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04-v1.0
# 保存镜像为压缩包，命名为 cuda11.8.0-cudnn8-devel-ubuntu22.04-v1.0.tar
docker save -o cuda11.8.0-cudnn8-devel-ubuntu22.04-v1.0.tar nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04-v1.0
# 迁移后加载压缩包到 docker image 中，参考 1 中命令照常启动
docker load -i cuda11.8.0-cudnn8-devel-ubuntu22.04-v1.0.tar

这样当前的训练环境被打包成了 tar 文件，便于后续的环境迁移。