Use torch generic workflow for CI

ghstack-source-id: 008612abbfe85656a1cae9ef590710687744bfd6
Pull Request resolved: #325

.github/workflows/unit_test_4gpu.yaml

@@ -9,34 +9,24 @@ concurrency:
   group: unit-test${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
   cancel-in-progress: true
 
-defaults:
-  run:
-    shell: bash -l -eo pipefail {0}
-
 jobs:
-  unit_tests_4gpu:
-    runs-on: linux.g5.12xlarge.nvidia.gpu
-    strategy:
-      matrix:
-        python-version: ['3.10']
-    steps:
-      - name: Check out repo
-        uses: actions/checkout@v3
-      - name: Setup conda env
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          auto-update-conda: true
-          miniconda-version: "latest"
-          activate-environment: test
-          python-version: ${{ matrix.python-version }}
-      - name: Update pip
-        run: python -m pip install --upgrade pip
-      - name: Install dependencies
-        run: |
-          pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
-          python -m pip install -r requirements.txt
-          python -m pip install -r dev-requirements.txt
-      - name: Run test_runner.py
-        run: python ./test_runner.py
-      - name: Upload Coverage to Codecov
-        uses: codecov/codecov-action@v3
+  build-test:
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.g5.12xlarge.nvidia.gpu
+      gpu-arch-type: cuda
+      gpu-arch-version: "12.1"
+      docker-image: "nvidia/cuda:12.4.1-runtime-ubuntu22.04"
+      repository: "pytorch/torchtitan"
+      upload-artifact: "outputs"
+      script: |
+        pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+        python -m pip install -r requirements.txt
+        python -m pip install -r dev-requirements.txt
+        python ./test_runner.py
+        mv outputs artifacts-to-be-uploaded
+
+# TODO the below were not ported over to the 'generic workflow' - do they matter?
+
+#       - name: Upload Coverage to Codecov
+#         uses: codecov/codecov-action@v3