build torch and apex in a base image

eole-nlp · Jun 4, 2024 · 66d1c21 · 66d1c21
1 parent 8917d33
commit 66d1c21
Show file tree

Hide file tree

Showing 7 changed files with 138 additions and 53 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -7,22 +7,28 @@ on:
  description: "EOLE version"
  required: true
  type: string
- # to facilitate initial tests in PR
- push:
- branches:
- - "docker"
+ torch_version:
+ description: "PyTorch version"
+ required: true
+ type: string
+ cuda_version:
+ description: "CUDA version"
+ required: true
+ type: string
 
-run-name: ${{ github.workflow }} -- ${{ inputs.eole_version || 'test' }}
+run-name: ${{ github.workflow }} -- ${{ inputs.eole_version }} torch:${{ inputs.torch_version }}, cuda:${{ inputs.cuda_version }}
 
 env:
- EOLE_VERSION: ${{ inputs.eole_version || 'test' }}
+ EOLE_VERSION: ${{ inputs.eole_version }}
+ TORCH_VERSION: ${{ inputs.torch_version }}
+ CUDA_VERSION: ${{ inputs.cuda_version }}
 
 jobs:
  build:
  runs-on: ubuntu-22.04
- strategy:
- matrix:
- cuda_version: [11.8.0, 12.1.0]
+ # strategy:
+ #  matrix:
+ #  cuda_version: [11.8.0, 12.1.0]
  permissions: write-all
  steps:
  - name: Checkout repo
@@ -35,4 +41,4 @@ jobs:
  password: ${{ secrets.GITHUB_TOKEN }}
  - name: Build
  run: |
- docker/build.sh ${{ env.EOLE_VERSION }} ${{ matrix.cuda_version}}
+ docker/build.sh ${{ env.EOLE_VERSION }} ${{ env.TORCH_VERSION}} ${{ env.CUDA_VERSION}}
diff --git a/.github/workflows/build_base.yml b/.github/workflows/build_base.yml
@@ -0,0 +1,39 @@
+name: Build Base Docker Image
+
+on:
+ workflow_dispatch:
+ inputs:
+ torch_version:
+ description: "PyTorch version"
+ required: true
+ type: string
+ cuda_version:
+ description: "CUDA version"
+ required: true
+ type: string
+
+run-name: ${{ github.workflow }} -- torch:${{ inputs.torch_version }} -- cuda:${{ inputs.cuda_version }}
+
+env:
+ TORCH_VERSION: ${{ inputs.torch_version }}
+ CUDA_VERSION: ${{ inputs.cuda_version }}
+
+jobs:
+ build:
+ runs-on: ubuntu-22.04
+ # strategy:
+ # matrix:
+ # cuda_version: [11.8.0, 12.1.0]
+ permissions: write-all
+ steps:
+ - name: Checkout repo
+ uses: actions/checkout@v4
+ - name: Login to ghcr
+ uses: docker/login-action@v3
+ with:
+ registry: ghcr.io
+ username: ${{ github.actor }}
+ password: ${{ secrets.GITHUB_TOKEN }}
+ - name: Build
+ run: |
+ docker/build_base.sh ${{ env.TORCH_VERSION }} ${{ env.CUDA_VERSION }}
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -1,40 +1,7 @@
+ARG TORCH_VERSION=2.3.0
 ARG CUDA_VERSION=12.1.0
-FROM nvidia/cuda:$CUDA_VERSION-devel-ubuntu22.04
 
-RUN apt-get update && apt-get install -y locales gcc g++ python3-dev
-RUN apt-get update && apt-get install -y \
- git \
- python3-pip \
- python3-dev \
- libprotobuf-dev \
- libprotobuf-c-dev
-
-RUN pip3 install --upgrade pip
-RUN pip3 install packaging
-
-# Install torch
-RUN CU=$(echo "${CUDA_VERSION%.*}" | sed 's/\.//g'); pip3 install torch --index-url "https://download.pytorch.org/whl/cu$CU"
-
-# Install apex
-RUN mkdir /setup
-WORKDIR /setup
-RUN git clone https://github.com/nvidia/apex
-WORKDIR /setup/apex
-RUN pip3 install ninja
-ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.6"
-RUN pip3 install -v --no-build-isolation \
- --config-settings --global-option="--cpp_ext" \
- --config-settings --global-option="--cuda_ext" \
- --config-settings --global-option="--deprecated_fused_adam" \
- --global-option="--xentropy" \
- --global-option="--fast_multihead_attn" \
- ./
-
-# Install flash-attention
-RUN pip install flash-attn --no-build-isolation
-
-# Install AutoAWQ
-RUN pip install autoawq
+FROM ghcr.io/eole-nlp/eole-base:torch$TORCH_VERSION-ubuntu22.04-cuda$CUDA_VERSION
 
 COPY . /eole
 WORKDIR /eole

diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base
@@ -0,0 +1,43 @@
+# base image with torch and apex
+# relatively lighter than the full-fledged ngc pytorch images
+# ARG TORCH_VERSION=2.3.0
+ARG CUDA_VERSION=12.1.0
+FROM nvidia/cuda:$CUDA_VERSION-devel-ubuntu22.04
+
+ARG TORCH_VERSION==2.3.0
+
+
+RUN apt-get update && apt-get install -y \
+ libprotobuf-dev \
+ libprotobuf-c-dev \
+ g++ \
+ gcc \
+ git \
+ locales \
+ python3-dev \
+ python3-pip
+
+
+RUN pip3 install --upgrade pip
+RUN pip3 install packaging
+
+# Install torch
+RUN CU=$(echo "${CUDA_VERSION%.*}" | sed 's/\.//g'); pip3 install torch==$TORCH_VERSION --index-url "https://download.pytorch.org/whl/cu$CU"
+
+# Install apex
+RUN mkdir /setup
+WORKDIR /setup
+RUN git clone https://github.com/nvidia/apex
+WORKDIR /setup/apex
+RUN pip3 install ninja
+ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.6;8.7;9.0+PTX"
+RUN export MAX_JOBS=$(nproc); pip3 install -v --no-build-isolation \
+ --config-settings --global-option="--cpp_ext" \
+ --config-settings --global-option="--cuda_ext" \
+ --config-settings --global-option="--deprecated_fused_adam" \
+ --global-option="--xentropy" \
+ --global-option="--fast_multihead_attn" \
+ ./
+
+# Install flash-attention
+RUN pip install flash-attn --no-build-isolation
diff --git a/docker/build.sh b/docker/build.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 #
-# Build and push version X of OpenNMT-py with CUDA Y:
-# ./build.sh X Y
+# Build and push version X of EOLE with torch Y and CUDA Z:
+# ./build.sh X Y Z
 
 set -e
 
@@ -15,14 +15,15 @@ DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 ROOT_DIR=$DIR/..
 cd $ROOT_DIR
 
+
 EOLE_VERSION="$1"
-CUDA_VERSION="$2"
-[ -z "$CUDA_VERSION" ] && CUDA_VERSION="12.1.0"
+TORCH_VERSION="$2"
+CUDA_VERSION="$3"
 
 IMAGE="ghcr.io/eole-nlp/eole"
-TAG="$EOLE_VERSION-ubuntu22.04-cuda${CUDA_VERSION%.*}"
+TAG="$EOLE_VERSION-torch$TORCH_VERSION-ubuntu22.04-cuda${CUDA_VERSION%.*}"
 
-echo "Building $IMAGE:$TAG with CUDA_VERSION=$CUDA_VERSION"
+echo "Building $IMAGE:$TAG with TORCH_VERSION=$TORCH_VERSION,CUDA_VERSION=$CUDA_VERSION"
 
-docker build -t $IMAGE:$TAG --progress=plain -f docker/Dockerfile --build-arg CUDA_VERSION=$CUDA_VERSION --no-cache .
+docker build -t $IMAGE:$TAG --progress=plain -f docker/Dockerfile --build-arg TORCH_VERSION=$TORCH_VERSION --build-arg CUDA_VERSION=$CUDA_VERSION --no-cache .
 docker push $IMAGE:$TAG
diff --git a/docker/build_base.sh b/docker/build_base.sh
@@ -0,0 +1,28 @@
+#!/bin/bash
+#
+# Build and push base image with torch version X and CUDA version Y:
+# ./build.sh X Y
+
+set -e
+
+# allow user to run this script from anywhere
+# from https://stackoverflow.com/a/246128
+# one-liner which will give you the full directory name
+# of the script no matter where it is being called from
+unset CDPATH
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+ROOT_DIR=$DIR/..
+cd $ROOT_DIR
+
+TORCH_VERSION="$1"
+CUDA_VERSION="$2"
+[ -z "$CUDA_VERSION" ] && CUDA_VERSION="12.1.0"
+
+IMAGE="ghcr.io/eole-nlp/eole-base"
+TAG="torch$TORCH_VERSION-ubuntu22.04-cuda$CUDA_VERSION"
+
+echo "Building $IMAGE:$TAG with TORCH_VERSION=$TORCH_VERSION,CUDA_VERSION=$CUDA_VERSION"
+
+docker build -t $IMAGE:$TAG --progress=plain -f docker/Dockerfile-base --build-arg TORCH_VERSION=$TORCH_VERSION --build-arg CUDA_VERSION=$CUDA_VERSION --no-cache .
+docker push $IMAGE:$TAG
diff --git a/requirements.opt.txt b/requirements.opt.txt
@@ -3,4 +3,5 @@ rapidfuzz
 scipy
 bitsandbytes>=0.41.2
 spacy
-gradio
+gradio
+autoawq