[CI] Fix benchmarks for LLMs

.github/workflows/benchmarks.yml

@@ -84,7 +84,7 @@ jobs:
           python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers
           python -m pip install "pybind11[global]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          python3.10 -m pip install safetensors tqdm pandas numpy matplotlib
+          python3.10 -m pip install safetensors tqdm pandas numpy matplotlib ray
           python3.10 setup.py develop
 
           # test import

.github/workflows/benchmarks_pr.yml

@@ -79,7 +79,7 @@ jobs:
           ${{ matrix.device == 'CPU' && 'export CUDA_VISIBLE_DEVICES=' || '' }}
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 -U
-          python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers
+          python3.10 -m pip install ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]" transformers ray
           python3.10 -m pip install "pybind11[global]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
           python3.10 -m pip install safetensors tqdm pandas numpy matplotlib

benchmarks/test_llm.py

@@ -16,6 +16,11 @@
 
 _has_transformers = importlib.import_module("transformers") is not None
 
+# Skip all these tests if gpu is not available
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="GPU not available"
+)
+
 
 @pytest.fixture(scope="module")
 def transformers_wrapper():