Add Megatron CI workflow (deepspeedai#2614)

mrwyattii · tjruwase · web-flow · commit df985fac41a7 · 2022-12-17T20:23:13.000-05:00
* added megatron unit test

* Update nv-megatron.yml

Co-authored-by: Olatunji Ruwase &lt;olruwase@microsoft.com&gt;
diff --git a/.github/workflows/nv-megatron.yml b/.github/workflows/nv-megatron.yml
@@ -0,0 +1,57 @@
+name: nv-megatron
+
+on:
+  push:
+    branches:
+      - 'master'
+      - 'staging**'
+    paths-ignore:
+      - 'docs/**'
+  pull_request:
+    paths-ignore:
+      - 'docs/**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    runs-on: [self-hosted, nvidia, cu116, v100]
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: environment
+        run: |
+          echo "JobID: $AISC_NODE_INSTANCE_ID"
+          nvidia-smi
+          which python
+          python --version
+          which nvcc
+          nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision triton
+          pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu116
+          python -c "import torch; print('torch:', torch.__version__, torch)"
+          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+      - name: Install
+        run: |
+          pip uninstall --yes deepspeed
+          pip install .[dev]
+          pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" git+https://github.com/NVIDIA/apex.git
+
+      - name: Python environment
+        run: |
+          pip list
+
+      - name: Unit tests
+        run: |
+          git clone --branch mrwyattii/add-unit-test https://github.com/microsoft/Megatron-DeepSpeed.git
+          cd Megatron-DeepSpeed
+          pip install .
+          unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
+          if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
+          cd tests
+          MEGATRON_CKPT_DIR=/blob/megatron_ckpt/ TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --verbose ./
