feat: run thunder tests as part of LitGPT CI (#1975)

Author: deependujha

.azure/gpu-test.yml

@@ -75,15 +75,25 @@ jobs:
         displayName: "Env details"
 
       - bash: |
-          pytest -v --disable-pytest-warnings --strict-markers --color=yes \
+          pytest -v \
             --ignore-glob="tests/test_thunder*" \
             --ignore="tests/test_unsloth_executor.py"
         displayName: "Ordinary tests"
         condition: ne(variables['dependency'], 'compiler')
         timeoutInMinutes: "5"
 
       - bash: |
-          pytest -v --disable-pytest-warnings --strict-markers --color=yes
+          # install thunder from source, so that, thunder.tests will be available
+          pip install -U "thunder[test] @ git+https://github.com/Lightning-AI/lightning-thunder.git"
+          PL_RUN_CUDA_TESTS=0 pytest tests/ext_thunder/test_thunder_networks.py -v # without env var, it filters out all tests
+        displayName: "Extra tests w. Thunder [main branch]"
+        condition: eq(variables['dependency'], 'compiler')
+        env:
+          PL_RUN_CUDA_TESTS: "0"
+        timeoutInMinutes: "10"
+
+      - bash: |
+          pytest -v
         displayName: "All tests"
         condition: eq(variables['dependency'], 'compiler')
         timeoutInMinutes: "5"

pyproject.toml

@@ -75,6 +75,7 @@ optional-dependencies.test = [
   "einops>=0.7",
   "protobuf>=4.23.4",
   "pytest>=8.1.1",
+  "pytest-benchmark>=5.1",
   "pytest-dependency>=0.6",
   "pytest-rerunfailures>=14",
   "pytest-timeout>=2.3.1",

tests/ext_thunder/test_thunder_networks.py

@@ -0,0 +1,8 @@
+"""Run thunder tests as part of LitGPT CI"""
+
+from litgpt.utils import _THUNDER_AVAILABLE
+
+if _THUNDER_AVAILABLE:
+    from thunder.tests.test_networks import *  # noqa: F403
+else:
+    print("Skipping test_thunder_networks.py (thunder not available)")