use gpu machine to build image

Azure · Feb 6, 2024 · 622e2ff · 622e2ff
1 parent 9e4f079
commit 622e2ff
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 8 deletions.
diff --git a/.github/workflows/build-image.yaml b/.github/workflows/build-image.yaml
@@ -14,7 +14,7 @@ on:
 jobs:
   docker:
     name: Docker build ${{ matrix.name }}
-    runs-on: [self-hosted, linux, x64, cpu]
+    runs-on: [self-hosted, linux, x64, gpu]
     timeout-minutes: 600
     permissions:
       contents: read
@@ -56,8 +56,6 @@ jobs:
           fi
           DOCKERFILE=dockerfile/${{ matrix.name }}.dockerfile
 
-          BUILD_ARGS="NUM_MAKE_JOBS=8"
-
           CACHE_FROM="type=registry,ref=$(cut -d, -f1 <<< ${TAGS})"
           CACHE_TO=""
           if [[ "${{ github.event_name }}" != "pull_request" ]]; then

diff --git a/dockerfile/torch1.14-cuda11.8.dockerfile b/dockerfile/torch1.14-cuda11.8.dockerfile
@@ -49,7 +49,6 @@ 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 git+https://github.com/NVIDIA/TransformerEngine.git@stable
 

diff --git a/dockerfile/torch2.1-cuda12.2.dockerfile b/dockerfile/torch2.1-cuda12.2.dockerfile
@@ -49,7 +49,6 @@ 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 git+https://github.com/NVIDIA/TransformerEngine.git@stable
 

diff --git a/msamp/te/modules.py b/msamp/te/modules.py
@@ -69,7 +69,7 @@ def set_fp8_weights(self):
                     self, weight_cast_attr,
                     Float8Tensor(
                         data=torch.empty(
-                            shape,
+                            (0, 0),
                             device=torch.cuda.current_device(),
                             dtype=torch.uint8,
                         ),
@@ -82,8 +82,7 @@ def set_fp8_weights(self):
                     self, weight_transpose_attr,
                     Float8Tensor(
                         data=torch.empty(
-                            shape[1],
-                            shape[0],
+                            (0, 0),
                             device=torch.cuda.current_device(),
                             dtype=torch.uint8,
                         ),