Based on extensive profiling:

- TensorRT 23.05
- oneAPI 2023.1.0
- oneDNN 3.1.1
- clang 14

Also include ctranslate2 fork with whisper-distil support and transformers 4.35.0

Dockerfile

@@ -1,8 +1,30 @@
+# A lot of software, a lot of versions
+
+ARG CLANG_VER=14
+ARG CTRANSLATE2_VER=cd26b3e
+ARG NVIDIA_VER=23.05
+ARG ONEAPI_VER=2023.1.0
+ARG ONEDNN_VER=3.1.1
+ARG TORCH_VER=2.1.0
+ARG TORCH_AUDIO_VER=2.1.0
+
+# Misc
+ARG CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;8.9;9.0+PTX"
+ARG CTRANSLATE2_ROOT="/opt/ctranslate2"
+ARG CTRANSLATE2_URL="https://github.com/chiiyeh/CTranslate2.git"
+
 # Builder
-FROM nvcr.io/nvidia/tensorrt:23.08-py3 as builder
+FROM nvcr.io/nvidia/tensorrt:${NVIDIA_VER}-py3 as builder
 
-# Set in environment in case we need to build any extensions
-ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;8.9;9.0+PTX"
+ARG CLANG_VER
+ARG CTRANSLATE2_ROOT
+ENV CTRANSLATE2_ROOT=${CTRANSLATE2_ROOT}
+ARG CTRANSLATE2_URL
+ARG CTRANSLATE2_VER
+ARG CUDA_ARCH_LIST
+ARG ONEAPI_VER
+ARG ONEDNN_VER
+ARG TORCH_CUDA_ARCH_LIST=${CUDA_ARCH_LIST}
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
@@ -11,60 +33,58 @@ RUN apt-get update && \
         wget \
         lsb-release \
         software-properties-common \
-        gnupg \
-        && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        gnupg
 
 WORKDIR /root
 
-# Install clang 16
+# Install clang
 RUN wget -O llvm.sh https://apt.llvm.org/llvm.sh && chmod +x llvm.sh
-RUN ./llvm.sh 16
+RUN ./llvm.sh ${CLANG_VER}
+
+RUN apt-get update && apt-get install -y --no-install-recommends libomp-${CLANG_VER}-dev
 
-# Use clang 16
-ENV CC=/usr/bin/clang-16
-ENV CXX=/usr/bin/clang++-16
+# Use clang
+ENV CC=/usr/bin/clang-${CLANG_VER}
+ENV CXX=/usr/bin/clang++-${CLANG_VER}
 
-ENV ONEAPI_VERSION=2023.2.0
+# Install oneAPI
 RUN wget -q https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB && \
     apt-key add *.PUB && \
     rm *.PUB && \
     echo "deb https://apt.repos.intel.com/oneapi all main" > /etc/apt/sources.list.d/oneAPI.list && \
-    apt-get update && \
-    apt-get install -y --no-install-recommends \
-        intel-oneapi-mkl-devel-$ONEAPI_VERSION \
+    apt-get update && apt-get install -y --no-install-recommends \
+        intel-oneapi-mkl-devel-${ONEAPI_VER} \
         && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
+# Install cmake
 RUN --mount=type=cache,target=/root/.cache pip install cmake==3.22.*
 
-RUN apt-get update && apt-get install -y libomp-16-dev
-
-ENV ONEDNN_VERSION=3.3
-RUN wget -q https://github.com/oneapi-src/oneDNN/archive/refs/tags/v${ONEDNN_VERSION}.tar.gz && \
+# Install oneDNN
+RUN wget -q https://github.com/oneapi-src/oneDNN/archive/refs/tags/v${ONEDNN_VER}.tar.gz && \
     tar xf *.tar.gz && \
     rm *.tar.gz && \
     cd oneDNN-* && \
-    cmake -DCMAKE_BUILD_TYPE=Release -DONEDNN_LIBRARY_TYPE=STATIC -DONEDNN_BUILD_EXAMPLES=OFF -DONEDNN_BUILD_TESTS=OFF -DONEDNN_ENABLE_WORKLOAD=INFERENCE -DONEDNN_ENABLE_PRIMITIVE="CONVOLUTION;REORDER" -DONEDNN_BUILD_GRAPH=OFF . && \
+    cmake -DCMAKE_BUILD_TYPE=Release -DONEDNN_LIBRARY_TYPE=STATIC -DONEDNN_BUILD_EXAMPLES=OFF \
+        -DONEDNN_BUILD_TESTS=OFF -DONEDNN_ENABLE_WORKLOAD=INFERENCE \
+        -DONEDNN_ENABLE_PRIMITIVE="CONVOLUTION;REORDER" -DONEDNN_BUILD_GRAPH=OFF . && \
     make -j$(nproc) install && \
     cd .. && \
     rm -r oneDNN-*
 
-RUN git clone --recursive https://github.com/OpenNMT/CTranslate2.git
+# Build CTranslate2
+RUN git clone --recursive ${CTRANSLATE2_URL}
 
 WORKDIR /root/CTranslate2
 
-RUN git checkout 2203ad5
+RUN git checkout ${CTRANSLATE2_VER}
 
 ARG CXX_FLAGS
 ENV CXX_FLAGS=${CXX_FLAGS:-"-msse4.1"}
 ARG CUDA_NVCC_FLAGS
 ENV CUDA_NVCC_FLAGS=${CUDA_NVCC_FLAGS:-"-Xfatbin=-compress-all"}
-ARG CUDA_ARCH_LIST
-ENV CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-"Common"}
-ENV CTRANSLATE2_ROOT=/opt/ctranslate2
+ENV CUDA_ARCH_LIST=${CUDA_ARCH_LIST:-"Common"}
 
 RUN mkdir build && \
     cd build && \
@@ -79,42 +99,47 @@ COPY README.md .
 
 RUN --mount=type=cache,target=/root/.cache cd python && \
     pip --no-cache-dir install -r install_requirements.txt && \
-    python3 setup.py bdist_wheel --dist-dir $CTRANSLATE2_ROOT
+    python3 setup.py bdist_wheel --dist-dir ${CTRANSLATE2_ROOT}
 
 # Runtime
+FROM nvcr.io/nvidia/tensorrt:${NVIDIA_VER}-py3
 
-FROM nvcr.io/nvidia/tensorrt:23.08-py3
+ARG CLANG_VER
+ARG CTRANSLATE2_ROOT
+ENV CTRANSLATE2_ROOT=${CTRANSLATE2_ROOT}
+ARG TORCH_VER
+ARG TORCH_AUDIO_VER
 
 WORKDIR /app
 
 # Install deps
-RUN apt-get update && apt-get install -y zstd git-lfs libsox3 gnupg software-properties-common && rm -rf /var/lib/apt/lists/*
+RUN apt-get update && apt-get install -y zstd git-lfs libsox3
 
-# Install clang deps - installs full clang, which we can remove later
-RUN wget -O llvm.sh https://apt.llvm.org/llvm.sh && chmod +x llvm.sh
-RUN ./llvm.sh 16 && rm llvm.sh
+# Install llvm repo
+RUN wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | apt-key add -
+RUN echo "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-${CLANG_VER} main" > /etc/apt/sources.list.d/llvm.list
 
-# Change to runtime later
-RUN apt-get update && apt-get install -y libomp-16-dev
+# Install clang omp runtime
+RUN apt-get update && apt-get install -y libomp5-${CLANG_VER} && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 # Install our torch ver matching cuda
-RUN --mount=type=cache,target=/root/.cache pip install torch==2.1.0 torchaudio==2.1.0
+RUN --mount=type=cache,target=/root/.cache pip install torch==${TORCH_VER} torchaudio==${TORCH_AUDIO_VER}
 
 COPY requirements.txt .
 # Run pip install with cache so we speedup subsequent rebuilds
 RUN --mount=type=cache,target=/root/.cache pip install -r requirements.txt
 
-# Install compiled ctranslate2
-ENV CTRANSLATE2_ROOT=/opt/ctranslate2
-
 # Dynamic library fixes
 # Workaround for speaker verification with torchaudio
 RUN ln -sf /usr/lib/x86_64-linux-gnu/libsox.so.3 /usr/lib/x86_64-linux-gnu/libsox.so
 RUN ldconfig
 
-# Set Ctranslate2 path
+# Set CTranslate2 path
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CTRANSLATE2_ROOT/lib
 
+# Install CTranslate2 from builder
 COPY --from=builder $CTRANSLATE2_ROOT $CTRANSLATE2_ROOT
 RUN python3 -m pip --no-cache-dir install $CTRANSLATE2_ROOT/*.whl && \
     rm $CTRANSLATE2_ROOT/*.whl

requirements.txt

@@ -81,7 +81,7 @@ sympy==1.12
 threadpoolctl==3.2.0
 tokenizers==0.14.1
 tqdm==4.66.1
-transformers==4.34.1
+transformers==4.35.0
 triton==2.1.0
 typing_extensions==4.7.1
 ujson==5.8.0