litgpt/.lightning/workflows/tests.yaml at 55521f2b8b643065a56506a0f986c563cfab7d0d · Lightning-AI/litgpt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
trigger:
  push:
    branches: ["main"]
  pull_request:
    branches: ["main"]

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
parametrize:
  matrix:
    dependency: ["", "compiler"]
  include: []
  exclude: []

env:
  SKIP_WITH_CI: "1" # skip single tests with CI
  NCCL_DEBUG: "INFO"
  CUBLAS_WORKSPACE_CONFIG: ":4096:8"
  NCCL_IGNORE_DISABLED_P2P: "1"
  TORCH_VERSION: "2.8.0"
  RUN_ONLY_CUDA_TESTS: "1" # run CUDA tests only

run: |
  whereis nvidia
  nvidia-smi
  python --version
  pip --version
  pip list
  set -ex

  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
    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

  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}'"

  pytest -v --durations=100

  if [ "${dependency}" == "" ]; then
    # Run the tokenizer tests (CPU-only, so excluded above) with the registry fallback for gated repos.
    RUN_ONLY_CUDA_TESTS=0 LITGPT_TOKENIZER_REGISTRY_FALLBACK=1 timeout 10m pytest tests/test_tokenizer.py -v -s -rs --timeout=100 --durations=50
  fi

  wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/run_standalone_tests.sh
  PL_RUN_STANDALONE_TESTS=1 bash run_standalone_tests.sh "tests"

  if [ "${dependency}" == "compiler" ]; then
    uv pip uninstall lightning-thunder transformers
    # install thunder from source, so that, thunder.tests will be available
    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
    TRANSFORMERS_VERSION=$(curl -fsSL https://raw.githubusercontent.com/Lightning-AI/lightning-thunder/main/requirements/test.txt \
      | grep '^transformers==' \
      | cut -d'=' -f3 \
      | cut -d'#' -f1 \
      | xargs)
    if [ -z "${TRANSFORMERS_VERSION}" ]; then
      echo "Error: Could not determine transformers version from lightning-thunder requirements"
      exit 1
    fi
    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