Fix/docker action (#3266)

* [docker] Add ARG VERSION to determine the Tag

* [workflow] fixed the version in the release docker workflow

---------
Co-authored-by: liuzeming <liuzeming@4paradigm.com>

.github/workflows/release_docker_after_merge.yml

@@ -26,7 +26,7 @@ jobs:
         run: |
           version=$(cat version.txt)
           tag=hpcaitech/colossalai:$version
-          docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 -t $tag ./docker
+	  docker build --build-arg http_proxy=http://172.17.0.1:7890 --build-arg https_proxy=http://172.17.0.1:7890 --build-arg VERSION=v${version} -t $tag ./docker
           echo "tag=${tag}" >> $GITHUB_OUTPUT
 
       - name: Log in to Docker Hub

docker/Dockerfile

@@ -15,7 +15,8 @@ RUN git clone https://github.com/NVIDIA/apex && \
     pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_layer_norm" ./
 
 # install colossalai
-RUN git clone https://github.com/hpcaitech/ColossalAI.git \
+ARG VERSION=1
+RUN git clone -b ${VERSION} https://github.com/hpcaitech/ColossalAI.git \
     && cd ./ColossalAI \
     && CUDA_EXT=1 pip install -v --no-cache-dir .