add fa2-related tests

.github/workflows/test_3d_parallelism.yaml → .github/workflows/3d_parallelism_unit_tests.yaml

@@ -64,4 +64,6 @@ jobs:
       run: pip freeze | tee installed.txt
 
     - name: Run tests
-      run: pytest --color=yes --durations=0 --verbose tests/
+      # NOTE: -m "not fa2" will run all the unit tests that don't have the mark
+      # "fa2" (these are FA2-related tests, we can't run it on T4)
+      run: pytest -m "not fa2" --color=yes --durations=0 --verbose tests/

.github/workflows/fa2_unit_tests.yaml

@@ -0,0 +1,64 @@
+name: Run FA2-related unit tests
+
+on:
+  workflow_dispatch:
+  push:
+    branches: [ main ]
+    # paths:
+    #   - "src/**/*.py"
+    #   - "examples/**/*.py"
+    #   - "tests/**/*.py"
+
+  pull_request:
+    branches: [ main ]
+    #paths:
+    #  - "src/**/*.py"
+    #  - "examples/**/*.py"
+    #  - "tests/**/*.py"
+
+jobs:
+  tests:
+    runs-on: [multi-gpu, nvidia-gpu, 8-t4, ci]
+    container:
+      image: runpod/pytorch:2.1.1-py3.10-cuda12.1.1-devel-ubuntu22.04
+      ports:
+        - 80
+      options: --gpus all --shm-size "8G"
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Python environment
+      run: |
+        which python
+        python --version
+
+    - name: Check Pytorch version
+      run: |
+        nvidia-smi
+        python -c "import torch; print('torch:', torch.__version__, torch)"
+        python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+    - name: Instal nanotron
+      run: |
+        python -m pip install --upgrade pip
+        pip install packaging
+        pip install wheel
+        pip install "flash-attn>=2.4.2"  --no-build-isolation
+        git clone https://github.com/huggingface/nanotron.git
+        cd nanotron
+        pip install -e .
+        pip install -e .[dev]
+        pip install -e .[test]
+
+    - name: Install test dependencies
+      run: |
+        pip install pytest
+        pip install pytest-cov
+
+    - name: Show installed libraries and their versions
+      run: pip freeze | tee installed.txt
+
+    - name: Run tests
+      # NOTE: -m fa2 will only run the unit tests that have the mark
+      # "fa2" (these are FA2-related tests)
+      run: pytest -m fa2 --color=yes --durations=0 --verbose tests/

tests/kernels/test_layer_norm.py

@@ -23,6 +23,7 @@
 
 
 # @pytest.mark.skipif(available_gpus() < 1, reason="Testing test_fused_layer_norm requires at least 1 gpus")
+@pytest.mark.fa2
 @pytest.mark.parametrize(
     "hidden_size",
     [1024, 1025],  # fused layer norm supports 1024 as hidden size but not 1025

tests/pytest.ini

@@ -1,2 +1,4 @@
 [pytest]
 addopts=-n 35
+markers =
+    fa2: FA2-related