@@ -0,0 +1,372 @@
# syntax=docker/dockerfile:1
# NOTE: Building this image requires Docker version >= 23.0.
#
# For reference:
# - https://docs.docker.com/build/dockerfile/frontend/#stable-channel
ARG CUDA_VERSION = 12 .1.0
FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-devel-ubuntu20.04
ARG DEBIAN_FRONTEND = "noninteractive"
ENV DEBIAN_FRONTEND = ${ DEBIAN_FRONTEND }
ENV MAMBA_ROOT_PREFIX = ~/micromamba
ARG ROOT_PASSWD = "root"
ENV ROOT_PASSWD = ${ ROOT_PASSWD }
WORKDIR /root
SHELL [ "/bin/bash" , "-c" ]
# Base tools
RUN <<EOT
add -apt-repository ppa:git-core/ppa -y
's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
's/#PasswordAuthentication yes/PasswordAuthentication yes/' /etc/ssh/sshd_config
's/PubkeyAuthentication no/PubkeyAuthentication yes/' /etc/ssh/sshd_config
's/^#Port 22/Port 22/' /etc/ssh/sshd_config
's/^Port [0-9]*/Port 22/' /etc/ssh/sshd_config
echo 'root:${ROOT_PASSWD}' | chpasswd
echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc
echo 1 | bash <( curl -s https://cdn.jsdelivr.net/gh/hotwa/MicroMamba_Installer@main/install.sh)
'EOF' >> ~/.bashrc
source ~/micromamba/etc/profile.d/micromamba.sh
alias mamba = micromamba
alias mba = mamba
# PyTorch
ARG CONDA_ENV_NAME = "deepspeed"
ENV CONDA_ENV_NAME = ${ CONDA_ENV_NAME }
ARG PYTHON_VERSION = 3 .10
ENV PYTHON_VERSION = ${ PYTHON_VERSION }
ENV PATH /opt/conda/bin:/opt/conda/envs/${ CONDA_ENV_NAME } /bin:$PATH
ARG PYTORCH_VERSION = 2 .3.0
ENV PYTORCH_VERSION = ${ PYTORCH_VERSION }
ENV TORCH_CUDA_ARCH_LIST = "7.0 7.2 7.5 8.0 8.6 8.7 8.9 9.0 9.0a"
ENV TORCH_NVCC_FLAGS = "-Xfatbin -compress-all"
ENV CMAKE_PREFIX_PATH = " $( dirname $( which conda) ) /../ "
ARG CUDA_NUM = '121'
ENV CUDA_NUM = ${ CUDA_NUM }
RUN <<EOT
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME } python = ${ PYTHON_VERSION } cmake ninja -c conda-forge -y
echo " conda activate ${ CONDA_ENV_NAME } " >> ~/.bashrc
${ CONDA_ENV_NAME }
# 获取指定版本的 PyTorch 源代码
${ STAGE_DIR } /pytorch
cd ${ STAGE_DIR } /pytorch
${ PYTORCH_VERSION }
# 安装依赖项
${ CUDA_NUM }
# 构建和安装 PyTorch
export CMAKE_PREFIX_PATH = ${ CONDA_PREFIX :- " $( dirname $( which conda) ) /../ " }
# 安装其他必要的依赖项
setuptools = = 69.5.1
[ deepspeed-testing]
torch = = ${ PYTORCH_VERSION } +${ CUDA } torchvision = = 0.18.0+${ CUDA } torchaudio = = 2.3.0 --extra-index-url https://download.pytorch.org/whl/${ CUDA }
= accelerate
# install apex
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
${ STAGE_DIR } /apex
cd apex
# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...
MAX_JOBS = 1 python3 -m pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
"import apex.amp; print('Apex is installed and the amp module is available.')"
cd ..
${ STAGE_DIR } /apex
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
# Pre-build **latest** DeepSpeed, so it would be ready for testing (otherwise, the 1st deepspeed test will timeout)
# This has to be run (again) inside the GPU VMs running the tests.
# The installation works here, but some tests fail, if we do not pre-build deepspeed again in the VMs running the tests.
# TODO: Find out why test fail. install deepspeed
# DS_BUILD_CPU_ADAM=${DS_BUILD_CPU_ADAM} DS_BUILD_FUSED_ADAM={DS_BUILD_FUSED_ADAM} python3 -m pip install "deepspeed<=0.14.0" --global-option="build_ext" --global-option="-j8" --no-cache -v --disable-pip-version-check 2>&1
# from https://github.com/huggingface/transformers/blob/main/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile install deepspeed fail
# reference deepspeed install from https://github.com/microsoft/DeepSpeed/blob/master/docker/Dockerfile
# install deepspeed prepare
# install Mellanox OFED
${ STAGE_DIR }
${ 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
cd ${ STAGE_DIR }
${ STAGE_DIR } /MLNX_OFED_LINUX-${ MLNX_OFED_VERSION } -ubuntu20.04-x86_64*
cd ..
# install nv_peer_mem
${ STAGE_DIR }
${ STAGE_DIR }
${ NV_PEER_MEM_TAG } ${ STAGE_DIR } /nv_peer_memory
cd ${ STAGE_DIR } /nv_peer_memory
cd ${ STAGE_DIR }
${ STAGE_DIR } /nvidia-peer-memory_${ NV_PEER_MEM_VERSION } .orig.tar.gz
cd ${ STAGE_DIR } /nvidia-peer-memory-${ NV_PEER_MEM_VERSION }
${ STAGE_DIR } /nvidia-peer-memory_${ NV_PEER_MEM_TAG } _all.deb
# install mpi
ENV PATH = /usr/local/mpi/bin:${ PATH }
ENV LD_LIBRARY_PATH = /usr/local/lib:/usr/local/mpi/lib:/usr/local/mpi/lib64:${ LD_LIBRARY_PATH }
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
# OPENMPI
${ STAGE_DIR }
${ STAGE_DIR }
cd ${ STAGE_DIR }
${ OPENMPI_BASEVERSION } /openmpi-${ OPENMPI_VERSION } .tar.gz | tar xzf -
cd openmpi-${ OPENMPI_VERSION }
= /usr/local/openmpi-${ OPENMPI_VERSION }
" $( nproc) " install
${ OPENMPI_VERSION } /usr/local/mpi
# Sanity check:
test -f /usr/local/mpi/bin/mpic++
cd ${ STAGE_DIR }
${ STAGE_DIR } /openmpi-${ OPENMPI_VERSION }
# Create a wrapper for OpenMPI to allow running as root by default
echo '#!/bin/bash' > /usr/local/mpi/bin/mpirun
echo 'mpirun.real --allow-run-as-root --prefix /usr/local/mpi "$@"' >> /usr/local/mpi/bin/mpirun
# Some Packages
RUN <<EOT
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
\
# SSH daemon port inside container cannot conflict with host OS port
ENV SSH_PORT = 2222
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
${ STAGE_DIR } /sshd_config && \
" 0,/^Port 22/s//Port ${ SSH_PORT } / " ${ STAGE_DIR } /sshd_config > /etc/ssh/sshd_config
# 29.78 Usage: install.sh [options...]
# 29.78
# 29.78 By default will install deepspeed and all third party dependencies across all machines listed in
# 29.78 hostfile (hostfile: /job/hostfile). If no hostfile exists, will only install locally
# 29.78
# 29.78 [optional]
# 29.78 -l, --local_only Install only on local machine
# 29.78 -s, --pip_sudo Run pip install with sudo (default: no sudo)
# 29.78 -r, --allow_sudo Allow script to be run by root (probably don't want this, instead use --pip_sudo)
# 29.78 -n, --no_clean Do not clean prior build state, by default prior build files are removed before building wheels
# 29.78 -m, --pip_mirror Use the specified pip mirror (default: the default pip mirror)
# 29.78 -H, --hostfile Path to MPI-style hostfile (default: /job/hostfile)
# 29.78 -e, --examples Checkout deepspeed example submodule (no install)
# 29.78 -v, --verbose Verbose logging
# 29.78 -h, --help This help text
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
1000 --shell /bin/bash deepspeed
echo "deepspeed ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers
# install cutlass https://github.com/NVIDIA/cutlass
# H100: architecture is Hopper (cutlass need add : cmake .. -DCUTLASS_NVCC_ARCHS="90a" )
# A100: architecture is Ampere
# V100: architecture is Volta
# T4: architecture is Turing
# ENV CUDACXX=${CUDA_INSTALL_PATH}/bin/nvcc
# 70： ）
# 75： ）
# 80： ）
# 90a： ）
# 89:GeForce RTX 4090
ARG DCUTLASS_NVCC_ARCHS = "89"
ENV DCUTLASS_NVCC_ARCHS = ${ DCUTLASS_NVCC_ARCHS }
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
cd /opt/cutlass
cd build
= ${ DCUTLASS_NVCC_ARCHS } -DCUTLASS_ENABLE_TESTS= OFF -DCUTLASS_UNITY_BUILD_ENABLED= ON # compiles for NVIDIA Hopper GPU architecture, like H100
" $( nproc) " install
cd ..
# make test_unit -j"$(nproc)"
# make test_unit_gemm_warp -j"$(nproc)"
# CUDA_ARCH_LIST="80;86;89;90"
ARG DEEPSPEED_INSTALL_FLAGS = "--allow_sudo --pip_sudo --no_clean"
ENV DEEPSPEED_INSTALL_FLAGS = ${ DEEPSPEED_INSTALL_FLAGS }
ARG CUDA_ARCH_LIST = "80;86;89;90"
ENV CUDA_ARCH_LIST = ${ CUDA_ARCH_LIST }
ARG DS_BUILD_SPARSE_ATTN = 0
ENV DS_BUILD_SPARSE_ATTN = ${ DS_BUILD_SPARSE_ATTN }
ARG DS_BUILD_FUSED_ADAM = 1
ENV DS_BUILD_FUSED_ADAM = ${ DS_BUILD_FUSED_ADAM }
ARG DS_BUILD_CPU_ADAM = 0
ENV DS_BUILD_CPU_ADAM = ${ DS_BUILD_CPU_ADAM }
ARG DS_BUILD_OPS = 1
ENV DS_BUILD_OPS = ${ DS_BUILD_OPS }
ENV CUTLASS_PATH = /opt/cutlass
# install deepspeed
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
setuptools = = ${ SETUPTOOLS_VERSION }
# install oneapi for deepspeed
${ STAGE_DIR } /oneCCL
cd ${ STAGE_DIR } /oneCCL
cd build
= /usr/local
${ STAGE_DIR } /DeepSpeed-Kernels
cd ${ STAGE_DIR } /DeepSpeed-Kernels
CUDA_ARCH_LIST = ${ CUDA_ARCH_LIST } python setup.py bdist_wheel
# pip install dist/deepspeed_kernels-*.whl
CUDA_ARCH_LIST = ${ CUDA_ARCH_LIST } pip install -v .
${ STAGE_DIR } /DeepSpeed
cd ${ STAGE_DIR } /DeepSpeed
DS_BUILD_OPS = ${ DS_BUILD_OPS } pip install dist/deepspeed*.whl --force-reinstall
# DS_BUILD_OPS=${DS_BUILD_OPS} pip install -r requirements/requirements.txt
# DS_BUILD_OPS=0 DS_BUILD_SPARSE_ATTN=0 DS_BUILD_CPU_ADAM=0 DS_BUILD_FUSED_ADAM=1 pip install -U --no-cache-dir .
# ./install.sh ${DEEPSPEED_INSTALL_FLAGS} --hostfile /job/hostfile # ./install.sh --allow_sudo --pip_sudo --no_clean --hostfile /path/to/your/hostfile
cd ..
# rm -rf ${STAGE_DIR}/DeepSpeed
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
# install transformers
${ STAGE_DIR } /transformers
cd ${ STAGE_DIR } /transformers
"pydantic<2"
# install flash-attn
# pip install packaging -i https://pypi.org/simple/ --trusted-host pypi.org
RUN <<EOT
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
${ CONDA_ENV_NAME }
\
\
\
\
# add vscode server
# RUN <<EOT
# #!/bin/bash
# wget -qO- https://update.code.visualstudio.com/commit:${commit_id}/server-linux-x64/stable
# code-server --install-extension ms-python.vscode-pylance
# EOT
# 启动 ssh 服务
# CMD ["/bin/bash", "-c", "service ssh start; tail -f /dev/null"]
CMD [ "/usr/sbin/sshd" , "-D" ]
