Updating to CUDA 12.8, pytorch 2.7 and deepseepd 0.17.2

docker/base/Dockerfile

@@ -1,6 +1,6 @@
 FROM continuumio/miniconda3:24.7.1-0 AS build
 
-ARG CUDA_VERSION=12.6.0
+ARG CUDA_VERSION=12.8.1
 ENV CUDA_VERSION=$CUDA_VERSION
 
 ARG PYTHON_VERSION=3.11.11
@@ -53,7 +53,8 @@ RUN <<EOR
       cuda-command-line-tools=${CUDA_VERSION} cuda-compiler=${CUDA_VERSION} cuda-runtime=${CUDA_VERSION} \
       cuda-libraries=${CUDA_VERSION} cuda-libraries-dev=${CUDA_VERSION} cuda-libraries-static=${CUDA_VERSION} \
       cuda-cudart=${CUDA_SHORT_VERSION} cuda-cudart_linux-64=${CUDA_SHORT_VERSION} \
-      cudnn=9.3 \
+      conda-forge::nccl=2.27.5.1 \
+      cudnn=9.10 \
       py-cpuinfo=9 \
       conda-forge::ffmpeg=7.1.0 \
       conda-forge::portaudio=19.7.0 \
@@ -76,7 +77,7 @@ RUN <<EOR
   conda activate alltalk
   mkdir -p ${ALLTALK_DIR}/pip_cache
   pip install --no-cache-dir --cache-dir=${ALLTALK_DIR}/pip_cache \
-      torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION_NO_DOT}
+      torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION_NO_DOT}
 
   if [ $? -ne 0 ] ; then
     echo "Failed to install pip dependencies: $RESULT"

docker/deepspeed/Dockerfile

@@ -2,7 +2,7 @@ ARG DOCKER_REPOSITORY
 ARG DOCKER_TAG=latest
 FROM ${DOCKER_REPOSITORY}alltalk_tts_environment:${DOCKER_TAG}
 
-ARG DEEPSPEED_VERSION=0.16.3
+ARG DEEPSPEED_VERSION=0.17.2
 ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
 
 WORKDIR ${ALLTALK_DIR}
@@ -27,15 +27,6 @@ RUN <<EOR
   cd ${ALLTALK_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/variables.sh

@@ -1,9 +1,9 @@
 #!/usr/bin/env bash
 
 declare -A ALLTALK_VARS
-ALLTALK_VARS["CUDA_VERSION"]=12.6.0
+ALLTALK_VARS["CUDA_VERSION"]=12.8.1
 ALLTALK_VARS["PYTHON_VERSION"]=3.11.11
-ALLTALK_VARS["DEEPSPEED_VERSION"]=0.16.2
+ALLTALK_VARS["DEEPSPEED_VERSION"]=0.17.2
 
 # Export single variables (needed by Docker locally)
 for key in "${!ALLTALK_VARS[@]}"