run on tpus

AI-Hypercomputer · Feb 26, 2024 · c97dcf9 · c97dcf9
1 parent e9fbd71
commit c97dcf9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/.github/workflows/BuildDockerImages.yml b/.github/workflows/BuildDockerImages.yml
@@ -28,6 +28,10 @@ on:
 
 jobs:
   build_images:
+    strategy:
+      fail-fast: false
+      matrix:
+        device-type: ["v4-8"]
     runs-on: ["self-hosted", "tpu", "${{ matrix.device-type }}"]
     steps:
     - uses: actions/checkout@v3

diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml
@@ -78,9 +78,9 @@ jobs:
       run: |
         docker run -v /home/runner/actions-runner/_work/maxtext/maxtext:/app --rm --privileged maxtext_base_image bash -c \
         'python3 MaxText/train.py MaxText/configs/base.yml run_name=runner_$(date +%Y-%m-%d-%H-%M) base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset steps=2 enable_checkpointing=false'
-    - name: Test build_and_upload.sh
-      run: |
-        bash build_and_upload.sh
+    # - name: Test build_and_upload.sh
+    #   run: |
+    #     bash build_and_upload.sh
     # - name: Test train.py with per_device_batch_size < 1
     #   run: |
     #     docker run -v /home/runner/actions-runner/_work/maxtext/maxtext:/app --rm --privileged maxtext_base_image bash -c \