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}