Adding docker build for DeepSpeed
This commit is contained in:
ekoc committed
1 parent
9a99126f68
commit
7522fa87d6
5 files changed
+250
-18
No files matched your search
@@ -0,0 +1,103 @@
|
||||
ARG CUDA_VERSION=12.1.1
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
|
||||
|
||||
ARG PYTHON_VERSION=3.11
|
||||
ENV PYTHON_VERSION=$PYTHON_VERSION
|
||||
|
||||
ARG PYTORCH_VERSION=2.2.1
|
||||
ENV PYTORCH_VERSION=$PYTORCH_VERSION
|
||||
|
||||
ARG DEEPSPEED_VERSION=0.16.1
|
||||
ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
##############################################################################
|
||||
# Directories:
|
||||
##############################################################################
|
||||
ENV STAGE_DIR=/tmp
|
||||
RUN mkdir -p ${STAGE_DIR}
|
||||
|
||||
##############################################################################
|
||||
# Installation/Basic Utilities
|
||||
##############################################################################
|
||||
SHELL ["/bin/bash", "-l", "-c"]
|
||||
ENV SHELL=/bin/bash
|
||||
|
||||
RUN <<EOR
|
||||
apt-get update
|
||||
apt-get install --no-install-recommends -y \
|
||||
software-properties-common \
|
||||
ca-certificates
|
||||
|
||||
add-apt-repository ppa:git-core/ppa -y # for latest git
|
||||
add-apt-repository ppa:deadsnakes/ppa -y # for python
|
||||
|
||||
apt-get update
|
||||
apt-get upgrade -y
|
||||
apt-get install --no-install-recommends -y \
|
||||
build-essential \
|
||||
autotools-dev \
|
||||
pdsh \
|
||||
cmake \
|
||||
g++ \
|
||||
gcc \
|
||||
curl \
|
||||
wget \
|
||||
vim \
|
||||
unzip \
|
||||
llvm-dev \
|
||||
git \
|
||||
python${PYTHON_VERSION}-dev \
|
||||
libcupti-dev \
|
||||
libaio-dev
|
||||
EOR
|
||||
|
||||
##############################################################################
|
||||
# Python & pip
|
||||
##############################################################################
|
||||
RUN <<EOR
|
||||
# Correct symlinks to use the proper python version:
|
||||
PYTHON_MAJOR_VERSION=${PYTHON_VERSION%%.*}
|
||||
rm -f /usr/bin/python${PYTHON_MAJOR_VERSION}
|
||||
ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_MAJOR_VERSION}
|
||||
ln -s /usr/bin/python3 /usr/bin/python
|
||||
|
||||
curl -O https://bootstrap.pypa.io/pip/3.7/get-pip.py
|
||||
python get-pip.py
|
||||
rm get-pip.py
|
||||
pip install --upgrade pip
|
||||
EOR
|
||||
|
||||
# Minimal dependencies needed to build deepspeed:
|
||||
RUN pip install \
|
||||
deepspeed-kernels \
|
||||
scikit-learn \
|
||||
torch==${PYTORCH_VERSION}
|
||||
|
||||
##############################################################################
|
||||
# DeepSpeed
|
||||
##############################################################################
|
||||
RUN <<EOR
|
||||
git clone https://github.com/microsoft/DeepSpeed.git ${STAGE_DIR}/DeepSpeed
|
||||
cd ${STAGE_DIR}/DeepSpeed
|
||||
git checkout .
|
||||
git checkout "tags/v${DEEPSPEED_VERSION}" -b "v${DEEPSPEED_VERSION}"
|
||||
EOR
|
||||
|
||||
##############################################################################
|
||||
# DeepSpeed build file
|
||||
##############################################################################
|
||||
RUN <<EOR
|
||||
cat << EOF > build_deepspeed.sh
|
||||
#!/usr/bin/env bash
|
||||
mkdir -p /deepspeed
|
||||
cd ${STAGE_DIR}/DeepSpeed
|
||||
DS_BUILD_OPS=1 python setup.py build_ext -j8 bdist_wheel
|
||||
mv ${STAGE_DIR}/DeepSpeed/dist/*.whl /deepspeed/
|
||||
EOF
|
||||
EOR
|
||||
|
||||
RUN chmod +x /build_deepspeed.sh
|
||||
|
||||
ENTRYPOINT ["/build_deepspeed.sh"]
|
||||
Executable
+65
@@ -0,0 +1,65 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
CUDA_VERSION=12.1.1
|
||||
PYTHON_VERSION=3.11
|
||||
PYTORCH_VERSION=2.2.1
|
||||
DEEPSPEED_VERSION=0.16.1
|
||||
|
||||
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
|
||||
cd $SCRIPT_DIR
|
||||
|
||||
# Parse arguments
|
||||
while [ "$#" -gt 0 ]; do
|
||||
case "$1" in
|
||||
--cuda-version)
|
||||
CUDA_VERSION="$2"
|
||||
shift
|
||||
;;
|
||||
--python-version)
|
||||
PYTHON_VERSION="$2"
|
||||
shift
|
||||
;;
|
||||
--pytorch-version)
|
||||
PYTORCH_VERSION="$2"
|
||||
shift
|
||||
;;
|
||||
--deepspeed-version)
|
||||
DEEPSPEED_VERSION="$2"
|
||||
shift
|
||||
;;
|
||||
*)
|
||||
# Allow to pass arbitrary arguments to docker as well to be flexible:
|
||||
echo "Unknown argument '$1'"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
PYTHON_VERSION_NO_DOT=${PYTHON_VERSION//./}
|
||||
if [[ -n $(find build -name "deepspeed-${DEEPSPEED_VERSION}*-cp${PYTHON_VERSION_NO_DOT}-cp${PYTHON_VERSION_NO_DOT}-*.whl") ]]
|
||||
then
|
||||
echo "DeepSpeed was already built - skipping..."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "Building DeepSpeed $DEEPSPEED_VERSION for CUDA $CUDA_VERSION using python ${PYTHON_VERSION} with PyTorch ${PYTORCH_VERSION}"
|
||||
|
||||
rm -rf build # make sure to properly clean up - we only want 1 wheel at the time
|
||||
mkdir -p build
|
||||
docker buildx \
|
||||
build \
|
||||
--build-arg CUDA_VERSION=$CUDA_VERSION \
|
||||
--build-arg PYTHON_VERSION=$PYTHON_VERSION \
|
||||
--build-arg PYTORCH_VERSION=$PYTORCH_VERSION \
|
||||
--build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
|
||||
-t deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION \
|
||||
.
|
||||
|
||||
docker run \
|
||||
--rm \
|
||||
-it \
|
||||
--gpus=all \
|
||||
--name deepspeed \
|
||||
-v $SCRIPT_DIR/build:/deepspeed \
|
||||
deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION
|
||||
Reference in new issue
Block a user