resolve failing tests with pt-2.1 (#21130)

* pip uninstall -y deepspeed * TORCH_VER * Apply suggestions from code review * A100 & L4 * machine

Author: Borda

.lightning/workflows/fabric.yml

@@ -4,20 +4,22 @@ trigger:
  pull_request:
  branches: ["master"]
 
-timeout: "75" # minutes
-machine: "L4_X_2"
+timeout: "55" # minutes
 parametrize:
  matrix: {}
  include:
- # note that this is setting also all oldest requirements which is linked to Torch == 2.0
+ # note that this is setting also all oldest requirements which is linked to Torch == 2.1
  - image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
  PACKAGE_NAME: "fabric"
- - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
+ machine: "A100_X_2"
+ - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
  PACKAGE_NAME: "fabric"
+ machine: "L4_X_2"
  # - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
  # PACKAGE_NAME: "fabric"
- - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
+ - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
  PACKAGE_NAME: "lightning"
+ machine: "L4_X_2"
  exclude: []
 
 env:
@@ -30,6 +32,7 @@ run: |
  python --version
  pip --version
  pip install -q fire wget packaging
+ pip list
  set -ex
 
  CUDA_VERSION="${image##*cuda}" # Remove everything up to and including "cuda"
@@ -40,12 +43,15 @@ run: |
  echo "Torch URL: ${TORCH_URL}"
  COVERAGE_SOURCE=$(python -c 'n = "$(PACKAGE_NAME)" ; print(dict(fabric="lightning_fabric").get(n, n))')
  echo "collecting coverage for: ${COVERAGE_SOURCE}"
+ TORCH_VER=$(python -c "import torch; print(torch.__version__.rsplit('.', 1)[0])")
 
  if [ "${TORCH_VER}" == "2.1" ]; then
  echo "Set oldest versions"
- cd requirements/fabric
+ pip uninstall -y deepspeed
  pip install -U "lightning-utilities[cli]"
+ cd requirements/fabric
  python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'strategies.txt']"
+ python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files strategies.txt
  cd ../..
  pip install "cython<3.0" wheel # for compatibility
  fi
@@ -92,6 +98,7 @@ run: |
  export PL_RUN_STANDALONE_TESTS=1
  wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/run_standalone_tests.sh
  bash ./run_standalone_tests.sh "tests_fabric"
+ export PL_RUN_STANDALONE_TESTS=0
 
  # echo "Reporting coverage" # todo
  # python -m coverage report

.lightning/workflows/pytorch.yml

@@ -4,20 +4,22 @@ trigger:
  pull_request:
  branches: ["master"]
 
-timeout: "75" # minutes
-machine: "L4_X_2"
+timeout: "55" # minutes
 parametrize:
  matrix: {}
  include:
- # note that this is setting also all oldest requirements which is linked to Torch == 2.0
+ # note that this is setting also all oldest requirements which is linked to Torch == 2.1
  - image: "pytorchlightning/pytorch_lightning:base-cuda12.1.1-py3.10-torch2.1"
  PACKAGE_NAME: "pytorch"
- - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
+ machine: "A100_X_2"
+ - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
  PACKAGE_NAME: "pytorch"
+ machine: "L4_X_2"
  # - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
  # PACKAGE_NAME: "pytorch"
- - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.7"
+ - image: "pytorchlightning/pytorch_lightning:base-cuda12.6.3-py3.12-torch2.8"
  PACKAGE_NAME: "lightning"
+ machine: "L4_X_2"
  exclude: []
 
 env:
@@ -30,6 +32,7 @@ run: |
  python --version
  pip --version
  pip install -q fire wget packaging
+ pip list
  set -ex
 
  CUDA_VERSION="${image##*cuda}" # Remove everything up to and including "cuda"
@@ -40,12 +43,15 @@ run: |
  echo "Torch URL: ${TORCH_URL}"
  COVERAGE_SOURCE=$(python -c 'n = "$(PACKAGE_NAME)" ; print(dict(fabric="pytorch_lightning").get(n, n))')
  echo "collecting coverage for: ${COVERAGE_SOURCE}"
+ TORCH_VER=$(python -c "import torch; print(torch.__version__.rsplit('.', 1)[0])")
 
  if [ "${TORCH_VER}" == "2.1" ]; then
- recho "Set oldest versions"
- cd requirements/pytorch
+ echo "Set oldest versions"
+ pip uninstall -y deepspeed
  pip install -U "lightning-utilities[cli]"
+ cd requirements/pytorch
  python -m lightning_utilities.cli requirements set-oldest --req_files "['base.txt', 'extra.txt', 'strategies.txt', 'examples.txt']"
+ python -m lightning_utilities.cli requirements prune-pkgs --packages deepspeed --req_files strategies.txt
  cd ../..
  pip install "cython<3.0" wheel # for compatibility
  fi
@@ -108,6 +114,7 @@ run: |
  export PL_RUN_STANDALONE_TESTS=1
  wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/run_standalone_tests.sh
  bash ./run_standalone_tests.sh "tests_pytorch"
+ export PL_RUN_STANDALONE_TESTS=0
 
  echo "Testing: PyTorch standalone tasks"
  cd tests_pytorch/

dockers/base-cuda/Dockerfile

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 ARG UBUNTU_VERSION=22.04
-ARG CUDA_VERSION=11.7.1
+ARG CUDA_VERSION=12.1.1
 
 
 FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}