add tensor parallel embedding test

axonn-ai · Jul 9, 2024 · 0d7b577 · 0d7b577
1 parent 4c3c079
commit 0d7b577
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/.github/workflows/nvidia-rtx-3090-tests.yaml b/.github/workflows/nvidia-rtx-3090-tests.yaml
@@ -49,6 +49,9 @@ jobs:
     - name: Run intra-layer Conv unit tests 
       run: |
         mpirun -mca orte_allowed_exit_without_sync 1 -n 2 pytest --with-mpi ./axonn/tests/test_intra_layer_conv.py 
+    - name: Run intra-layer Embedding unit tests
+        mpirun -mca orte_allowed_exit_without_sync 1 -n 2 pytest --with-mpi ./axonn/tests/test_intra_layer_emb.py -k bw_pass
+        mpirun -mca orte_allowed_exit_without_sync 1 -n 2 pytest --with-mpi ./axonn/tests/test_intra_layer_emb.py -k fw_pass
     - name: Uninstall AxoNN
       run: |
         pip uninstall --yes axonn