change dockerfile

Azure · Feb 4, 2024 · 1dac840 · 1dac840
1 parent 2150f64
commit 1dac840
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/.github/workflows/build-image.yaml b/.github/workflows/build-image.yaml
@@ -23,8 +23,8 @@ jobs:
         include:
         - name: torch1.14-cuda11.8
           tags: ghcr.io/azure/msamp:main-cuda11.8
-        - name: torch2.1-cuda12.1
-          tags: ghcr.io/azure/msamp:main-cuda12.1,ghcr.io/azure/msamp:latest
+        - name: torch2.1-cuda12.2
+          tags: ghcr.io/azure/msamp:main-cuda12.2,ghcr.io/azure/msamp:latest
     steps:
       - name: Checkout
         uses: actions/checkout@v2

diff --git a/dockerfile/torch2.1-cuda12.1.dockerfile → dockerfile/torch2.1-cuda12.2.dockerfile b/dockerfile/torch2.1-cuda12.1.dockerfile → dockerfile/torch2.1-cuda12.2.dockerfile
@@ -1,9 +1,9 @@
-FROM nvcr.io/nvidia/pytorch:23.04-py3
+FROM nvcr.io/nvidia/pytorch:23.10-py3
 
-# Ubuntu: 20.04
+# Ubuntu: 22.04
 # Python: 3.8
-# CUDA: 12.1.0
-# cuDNN: 8.9.0
+# CUDA: 12.2.0
+# cuDNN: 8.9.5
 # NCCL: v2.16.2-1 + FP8 Support
 # PyTorch: 2.1.0a0+fe05266f
 
@@ -48,9 +48,8 @@ RUN cd third_party/msccl && \
     -gencode=arch=compute_90,code=sm_90" && \
     make install
 # cache TE build to save time in CI
-ENV MAX_JOBS=1
 RUN python3 -m pip install --upgrade pip && \
-    python3 -m pip install flash-attn==1.0.9 git+https://github.com/NVIDIA/TransformerEngine.git@v0.11
+    python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
 ADD . .
 RUN python3 -m pip install . && \