Refactoring docker images, updating CUDA to 12.6 and pytorch 2.6

docker/deepspeed/Dockerfile

@@ -1,53 +1,15 @@
-FROM continuumio/miniconda3:24.7.1-0
+ARG GITHUB_REPOSITORY
+FROM ${GITHUB_REPOSITORY}alltalk_environment:latest
 
-ARG DEEPSPEED_VERSION=0.16.2
+ARG DEEPSPEED_VERSION=0.16.3
 ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
 
-SHELL ["/bin/bash", "-l", "-c"]
-ENV SHELL=/bin/bash
-ENV HOST=0.0.0.0
-ENV DEBIAN_FRONTEND=noninteractive
-ENV CUDA_DOCKER_ARCH=all
-ENV NVIDIA_VISIBLE_DEVICES=all
-ENV CONDA_AUTO_UPDATE_CONDA="false"
-
 ##############################################################################
 # Directories:
 ##############################################################################
 ENV STAGE_DIR=/tmp
 RUN mkdir -p ${STAGE_DIR}
 
-##############################################################################
-# Installation/Basic Utilities
-##############################################################################
-RUN <<EOR
-  apt-get update
-  apt-get upgrade -y
-  apt-get install --no-install-recommends -y \
-      espeak-ng \
-      curl \
-      wget \
-      jq \
-      vim
-
-   apt-get clean && rm -rf /var/lib/apt/lists/*
-EOR
-
-##############################################################################
-# Create a conda environment and install dependencies:
-##############################################################################
-COPY build/environment-*.yml environment.yml
-RUN <<EOR
-    RESULT=$( { conda env create -f environment.yml ; } 2>&1 )
-
-    if [ $? -ne 0 ] ; then
-      echo "Failed to install conda dependencies: $RESULT"
-      exit 1
-    fi
-
-    conda clean -a && pip cache purge
-EOR
-
 ##############################################################################
 # Install python dependencies
 ##############################################################################
@@ -56,8 +18,7 @@ RUN <<EOR
   conda activate alltalk
 
   pip install \
-      deepspeed-kernels \
-      scikit-learn
+      deepspeed-kernels
 EOR
 
 
@@ -69,6 +30,15 @@ RUN <<EOR
   cd ${STAGE_DIR}/DeepSpeed
   git checkout .
   git checkout "tags/v${DEEPSPEED_VERSION}" -b "v${DEEPSPEED_VERSION}"
+
+  # Patching some files due to bug https://github.com/microsoft/DeepSpeed/issues/6709
+  START_LINE=`awk '/#ifndef BF16_AVAILABLE/{ print NR; exit }' csrc/transformer/inference/csrc/gelu.cu`
+  TO_LINE=$((START_LINE + 2))
+  sed -i "${START_LINE},${TO_LINE}d" ./csrc/transformer/inference/csrc/gelu.cu
+
+  START_LINE=`awk '/#ifndef BF16_AVAILABLE/{ print NR; exit }' csrc/transformer/inference/csrc/transform.cu`
+  TO_LINE=$((START_LINE + 2))
+  sed -i "${START_LINE},${TO_LINE}d" ./csrc/transformer/inference/csrc/transform.cu
 EOR
 
 ##############################################################################

docker/deepspeed/build-deepspeed.sh

@@ -5,6 +5,8 @@ cd $SCRIPT_DIR
 
 . ${SCRIPT_DIR=}/../variables.sh
 
+DOCKER_TAG=latest
+
 # Parse arguments
 while [ "$#" -gt 0 ]; do
   case "$1" in
@@ -16,8 +18,17 @@ while [ "$#" -gt 0 ]; do
       DEEPSPEED_VERSION="$2"
       shift
       ;;
+    --tag)
+      DOCKER_TAG="$2"
+      shift
+      ;;
+    --github-repository)
+      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+        GITHUB_REPOSITORY="$2"
+        shift
+      fi
+      ;;
     *)
-      # Allow to pass arbitrary arguments to docker as well to be flexible:
       echo "Unknown argument '$1'"
       exit 1
       ;;
@@ -33,22 +44,16 @@ then
    exit 0
 fi
 
-CONDA_ENV=${SCRIPT_DIR=}/../conda/build/environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml
-if [ ! -f ${CONDA_ENV} ]; then
-    echo "No conda environment found. Please run 'build-conda-env.sh' first."
-    exit 1
-fi
-
 echo "Building DeepSpeed $DEEPSPEED_VERSION using python ${PYTHON_VERSION}"
 
 rm -rf build # make sure to properly clean up
 mkdir -p build
-cp ${CONDA_ENV} ${SCRIPT_DIR=}/build
 
 docker buildx \
   build \
+  --build-arg GITHUB_REPOSITORY=$GITHUB_REPOSITORY \
   --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
-  -t deepspeed:$DEEPSPEED_VERSION \
+  -t ${GITHUB_REPOSITORY}alltalk_deepspeed:${DOCKER_TAG} \
   .
 
 docker run \
@@ -57,4 +62,4 @@ docker run \
   --gpus=all \
   --name deepspeed \
   -v $SCRIPT_DIR/build:/deepspeed \
-  deepspeed:$DEEPSPEED_VERSION
+  ${GITHUB_REPOSITORY}alltalk_deepspeed:${DOCKER_TAG}