chore(workflows): Sync GPU CI workflow image with lightning-thunder repo (#2208)

This commit is contained in:
Bhimraj Yadav
2026-03-04 11:45:09 +00:00
committed by GitHub
parent e684dfb05b
commit 377951004c
+16 -8
View File
@@ -4,7 +4,7 @@ trigger:
pull_request:
branches: ["main"]
image: "pytorchlightning/lightning-thunder:ubuntu24.04-cuda12.6.3-cudnn-fe1.10.0-py3.10-pt_2.8.0-dev"
image: "pytorchlightning/lightning-thunder:ubuntu24.04-cuda12.8.1-cudnn-fe1.15.0-py3.12-pt_2.8.0-dev"
machine: "L4_X_2"
interruptible: "true"
timeout: "45" # minutes
@@ -30,16 +30,24 @@ run: |
pip list
set -ex
pip install -q '.[extra,test]' "torch==${TORCH_VERSION}" cffi -U --upgrade-strategy eager
echo "Install uv and create virtual environment"
curl -LsSf https://astral.sh/uv/install.sh | sh
[ -f "$HOME/.local/bin/env" ] && . "$HOME/.local/bin/env"
export PATH="$HOME/.local/bin:$PATH"
uv venv .venv --system-site-packages
. .venv/bin/activate
hash -r
uv pip install -q '.[extra,test]' "torch==${TORCH_VERSION}" cffi -U
if [ "${dependency}" == "compiler" ]; then
pip uninstall -y torchvision torchaudio
pip install -q '.[compiler,extra,test]' "torch==${TORCH_VERSION}"
uv pip uninstall torchvision torchaudio
uv pip install -q '.[compiler,extra,test]' "torch==${TORCH_VERSION}"
python -c "from thunder.executors import nvfuser_available ; assert nvfuser_available(), 'nvFuser is missing!'"
python -c "from thunder.executors.triton_utils import triton_version ; assert triton_version() is not None, 'triton is missing!'"
fi
pip list
uv pip list
python -c "import torch ; gpus = torch.cuda.device_count() ; assert gpus >= 2, f'GPU: {gpus}'"
python -c "from torch import __version__ as ver ; assert str(ver).split('+')[0] == '${TORCH_VERSION}', f'PyTorch: installed {ver} but expected ${TORCH_VERSION}'"
@@ -49,9 +57,9 @@ run: |
PL_RUN_STANDALONE_TESTS=1 bash run_standalone_tests.sh "tests"
if [ "${dependency}" == "compiler" ]; then
pip uninstall -y lightning-thunder transformers
uv pip uninstall lightning-thunder transformers
# install thunder from source, so that, thunder.tests will be available
pip install -U "lightning-thunder[test] @ git+https://github.com/Lightning-AI/lightning-thunder.git" "torch==${TORCH_VERSION}"
uv pip install -U "lightning-thunder[test] @ git+https://github.com/Lightning-AI/lightning-thunder.git" "torch==${TORCH_VERSION}"
# Pin transformers to match thunder's test_networks.py requirements
# See: https://github.com/Lightning-AI/lightning-thunder/blob/main/requirements/test.txt
# Get transformers version from thunder requirements
@@ -64,7 +72,7 @@ run: |
echo "Error: Could not determine transformers version from lightning-thunder requirements"
exit 1
fi
pip install transformers==${TRANSFORMERS_VERSION}
uv pip install transformers==${TRANSFORMERS_VERSION}
# without env var, it filters out all tests
RUN_ONLY_CUDA_TESTS=0 pytest tests/ext_thunder/test_thunder_networks.py -v
fi