You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

12 KiB

🐳 Dockerfile

Если вы планируете использовать собственный Dockerfile для запуска DeepSpeed (или любого распределенного обучения через SSH), не забудьте добавить в образ настройку pdsh, openssh и ca-certificates. Также полезно сразу указать sshd_config с настройками, чтобы SSH соединения оставались активными и использовали нестандартный порт.

🔧 Минимальный пример для добавления SSH в ваш Dockerfile

RUN apt-get update && \
    apt-get install -y pdsh ca-certificates openssh-client openssh-server

# Поддержка keep-alive для ssh
RUN echo "ClientAliveInterval 30" >> /etc/ssh/sshd_config

# Настройка порта для SSH
ENV SSH_PORT=2229
RUN sed -i "0,/^#Port 22/s//Port ${SSH_PORT}/" /etc/ssh/sshd_config

# Создание необходимых директорий и запуск sshd
RUN mkdir -p /run/sshd

CMD ["/usr/sbin/sshd", "-D"]

🐋 Полный пример Dockerfile для CUDA + SSH + DeepSpeed

Если хотите, можете сразу использовать готовую основу, которая включает CUDA, cuDNN и SSH-сервер для распределенного обучения:

FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu20.04

ENV DEBIAN_FRONTEND noninteractive

##############################################################################
# Temporary Installation Directory
##############################################################################
ENV STAGE_DIR=/tmp
RUN mkdir -p ${STAGE_DIR}

##############################################################################
# Installation/Basic Utilities
##############################################################################
RUN apt-get update && \
        apt-get install -y --no-install-recommends \
        software-properties-common build-essential autotools-dev \
        nfs-common pdsh \
        cmake g++ gcc \
        curl wget vim tmux emacs less unzip \
        htop iftop iotop ca-certificates openssh-client openssh-server \
        rsync iputils-ping net-tools sudo \
        llvm-dev

##############################################################################
# Installation Latest Git
##############################################################################
RUN add-apt-repository ppa:git-core/ppa -y && \
        apt-get update && \
        apt-get install -y git && \
        git --version

##############################################################################
# Client Liveness & Uncomment Port 22 for SSH Daemon
##############################################################################
# Keep SSH client alive from server side
RUN echo "ClientAliveInterval 30" >> /etc/ssh/sshd_config
RUN cp /etc/ssh/sshd_config ${STAGE_DIR}/sshd_config && \
        sed "0,/^#Port 22/s//Port 22/" ${STAGE_DIR}/sshd_config > /etc/ssh/sshd_config

##############################################################################
# Mellanox OFED
##############################################################################
ENV MLNX_OFED_VERSION=4.9-7.1.0.0
RUN apt-get install -y libnuma-dev
RUN cd ${STAGE_DIR} && \
        wget -q -O - http://www.mellanox.com/downloads/ofed/MLNX_OFED-${MLNX_OFED_VERSION}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64.tgz | tar xzf - && \
        cd MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64 && \
        ./mlnxofedinstall --user-space-only --without-fw-update --all -q && \
        cd ${STAGE_DIR} && \
        rm -rf ${STAGE_DIR}/MLNX_OFED_LINUX-${MLNX_OFED_VERSION}-ubuntu20.04-x86_64*

##############################################################################
# nv_peer_mem
##############################################################################
ENV NV_PEER_MEM_VERSION=1.2
ENV NV_PEER_MEM_TAG=${NV_PEER_MEM_VERSION}-0
RUN mkdir -p ${STAGE_DIR} && \
        git clone https://github.com/Mellanox/nv_peer_memory.git --branch ${NV_PEER_MEM_TAG} ${STAGE_DIR}/nv_peer_memory && \
        cd ${STAGE_DIR}/nv_peer_memory && \
        ./build_module.sh && \
        cd ${STAGE_DIR} && \
        tar xzf ${STAGE_DIR}/nvidia-peer-memory_${NV_PEER_MEM_VERSION}.orig.tar.gz && \
        cd ${STAGE_DIR}/nvidia-peer-memory-${NV_PEER_MEM_VERSION} && \
        apt-get update && \
        apt-get install -y dkms && \
        dpkg-buildpackage -us -uc && \
        dpkg -i ${STAGE_DIR}/nvidia-peer-memory_${NV_PEER_MEM_TAG}_all.deb

##############################################################################
# OPENMPI
##############################################################################
ENV OPENMPI_BASEVERSION=4.1
ENV OPENMPI_VERSION=${OPENMPI_BASEVERSION}.6
RUN cd ${STAGE_DIR} && \
        wget -q -O - https://download.open-mpi.org/release/open-mpi/v${OPENMPI_BASEVERSION}/openmpi-${OPENMPI_VERSION}.tar.gz | tar xzf - && \
        cd openmpi-${OPENMPI_VERSION} && \
        ./configure --prefix=/usr/local/openmpi-${OPENMPI_VERSION} && \
        make -j"$(nproc)" install && \
        ln -s /usr/local/openmpi-${OPENMPI_VERSION} /usr/local/mpi && \
        # Sanity check:
        test -f /usr/local/mpi/bin/mpic++ && \
        cd ${STAGE_DIR} && \
        rm -r ${STAGE_DIR}/openmpi-${OPENMPI_VERSION}
ENV PATH=/usr/local/mpi/bin:${PATH} \
        LD_LIBRARY_PATH=/usr/local/lib:/usr/local/mpi/lib:/usr/local/mpi/lib64:${LD_LIBRARY_PATH}
# Create a wrapper for OpenMPI to allow running as root by default
RUN mv /usr/local/mpi/bin/mpirun /usr/local/mpi/bin/mpirun.real && \
        echo '#!/bin/bash' > /usr/local/mpi/bin/mpirun && \
        echo 'mpirun.real --allow-run-as-root --prefix /usr/local/mpi "$@"' >> /usr/local/mpi/bin/mpirun && \
        chmod a+x /usr/local/mpi/bin/mpirun

##############################################################################
# Python
##############################################################################
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHON_VERSION=3
RUN apt-get install -y python3 python3-dev && \
        rm -f /usr/bin/python && \
        ln -s /usr/bin/python3 /usr/bin/python && \
        curl -O https://bootstrap.pypa.io/pip/3.6/get-pip.py && \
        python get-pip.py && \
        rm get-pip.py && \
        pip install --upgrade pip && \
        # Print python an pip version
        python -V && pip -V
RUN pip install pyyaml ipython build

##############################################################################
# Some Packages
##############################################################################
RUN apt-get update && \
        apt-get install -y --no-install-recommends \
        libsndfile-dev \
        libcupti-dev \
        libjpeg-dev \
        libpng-dev \
        screen \
        libaio-dev
RUN pip install psutil \
        yappi \
        cffi \
        ipdb \
        pandas \
        matplotlib \
        py3nvml \
        pyarrow \
        graphviz \
        astor \
        boto3 \
        tqdm \
        sentencepiece \
        msgpack \
        requests \
        sphinx \
        sphinx_rtd_theme \
        scipy \
        numpy \
        scikit-learn \
        nvidia-ml-py3 \
        mpi4py

##############################################################################
## SSH daemon port inside container cannot conflict with host OS port
###############################################################################
ENV SSH_PORT=2229
RUN cat /etc/ssh/sshd_config > ${STAGE_DIR}/sshd_config && \
        sed "0,/^Port 22/s//Port ${SSH_PORT}/" ${STAGE_DIR}/sshd_config > /etc/ssh/sshd_config

##############################################################################
# PyTorch
##############################################################################
RUN pip3 install torch==2.4.1+cu124 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu124 
##############################################################################
# PyYAML build issue
# https://stackoverflow.com/a/53926898
##############################################################################
RUN rm -rf /usr/lib/python3/dist-packages/yaml && \
        rm -rf /usr/lib/python3/dist-packages/PyYAML-*

##############################################################################
# DeepSpeed
##############################################################################
RUN git clone https://github.com/deepspeedai/DeepSpeed.git ${STAGE_DIR}/DeepSpeed
RUN cd ${STAGE_DIR}/DeepSpeed && \
        git checkout . && \
        git checkout master && \
        ./install.sh --pip_sudo --allow_sudo
RUN rm -rf ${STAGE_DIR}/DeepSpeed
RUN python -c "import deepspeed; print(deepspeed.__version__)"


WORKDIR /workspace

# Launch ssh server
RUN mkdir -p /run/sshd
CMD ["/usr/sbin/sshd", "-D"]

🔨 build.sh

docker build -t deepspeed_image_cuda_12_4 .

🚀 launch.sh

Перед запуском контейнера необходимо:

Кроме того, создаём пустую папку volume, в которой будут храниться файлы для DeepSpeed:

mkdir /media/Data1/common/docker/deepspeed/volume

📝 Скрипт launch.sh

#!/bin/bash

# Путь к директории проекта
DEEPSPEED_DIR="/media/Data1/common/docker/deepspeed"

# Имя контейнера
CONTAINER_NAME="deepspeed_container"

# Имя Docker-образа
IMAGE_NAME="deepspeed_image_cuda_12_4"

# Запуск контейнера
docker run --rm -it -d \
  -w /workspace \
  --name "$CONTAINER_NAME" \
  --network host \
  --shm-size=32g \
  --runtime nvidia --gpus '"device=1"' \
  -v "${DEEPSPEED_DIR}/volume":/workspace \
  -v "${DEEPSPEED_DIR}/ssh_docker":/root/.ssh \
  -e NCCL_SOCKET_IFNAME=10gb0,10gb1 \
 "$IMAGE_NAME"

🔍 Что здесь происходит

Опция Описание
-w /workspace Рабочая директория внутри контейнера.
--name "$CONTAINER_NAME" Имя контейнера для удобства управления.
--network host Используем сеть хоста (важно для скорости и для NCCL).
--shm-size=32g Размер памяти shm для PyTorch / NCCL.
--runtime nvidia --gpus ... Указываем Docker использовать конкретный GPU (например device=1 или all).
-v ${DEEPSPEED_DIR}/volume:/workspace Монтируем локальную папку для данных и логов в контейнер.
-v ${DEEPSPEED_DIR}/ssh_docker:/root/.ssh Монтируем SSH-конфиг и ключи для запуска multinode через SSH.
-e NCCL_SOCKET_IFNAME=10gb0,10gb1 Указываем интерфейсы для NCCL (10Gb).

🔌 connect.sh

docker exec -it deepspeed_container bash

После создания этих скриптов, не забудь сделать их исполняемыми:

chmod +x *.sh