docker: build images for latest PT 2.8 (#21042)

* docker: build images for latest PT `2.8`
* build release

(cherry picked from commit 20960ec)

Author: Borda

.github/checkgroup.yml

@@ -135,7 +135,8 @@ subprojects:
       - "build-pl (3.11, 2.4, 12.1.1)"
       - "build-pl (3.12, 2.5, 12.1.1)"
       - "build-pl (3.12, 2.6, 12.4.1)"
-      - "build-pl (3.12, 2.7, 12.6.3, true)"
+      - "build-pl (3.12, 2.7, 12.6.3)"
+      - "build-pl (3.12, 2.8, 12.6.3, true)"
 
   # SECTION: lightning_fabric
 

.github/workflows/docker-build.yml

@@ -49,7 +49,8 @@ jobs:
           - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.6", cuda_version: "12.4.1" }
-          - { python_version: "3.12", pytorch_version: "2.7", cuda_version: "12.6.3", latest: "true" }
+          - { python_version: "3.12", pytorch_version: "2.7", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.8", cuda_version: "12.6.3", latest: "true" }
     steps:
       - uses: actions/checkout@v4
         with:
@@ -97,7 +98,7 @@ jobs:
         # adding some more images as Thunder mainly using python 3.10,
         # and we need to support integrations as for example LitGPT
         python_version: ["3.10"]
-        pytorch_version: ["2.6.0", "2.7.1"]
+        pytorch_version: ["2.7.1", "2.8.0"]
         cuda_version: ["12.6.3"]
         include:
           # These are the base images for PL release docker images.
@@ -109,6 +110,7 @@ jobs:
           - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.6.0", cuda_version: "12.4.1" }
           - { python_version: "3.12", pytorch_version: "2.7.1", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.8.0", cuda_version: "12.6.3" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3
@@ -129,6 +131,7 @@ jobs:
             PYTHON_VERSION=${{ matrix.python_version }}
             PYTORCH_VERSION=${{ matrix.pytorch_version }}
             CUDA_VERSION=${{ matrix.cuda_version }}
+            MAKE_FLAGS="-j2"
           file: dockers/base-cuda/Dockerfile
           push: ${{ env.PUSH_NIGHTLY }}
           tags: "pytorchlightning/pytorch_lightning:base-cuda-py${{ matrix.python_version }}-torch${{ env.PT_VERSION }}-cuda${{ matrix.cuda_version }}"
@@ -157,6 +160,8 @@ jobs:
         continue-on-error: true
         uses: docker/build-push-action@v6
         with:
+          build-args: |
+            PYTORCH_VERSION="25.04"
           file: dockers/nvidia/Dockerfile
           push: false
         timeout-minutes: 55

dockers/base-cuda/Dockerfile

@@ -19,8 +19,9 @@ ARG CUDA_VERSION=11.7.1
 FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}
 
 ARG PYTHON_VERSION=3.10
-ARG PYTORCH_VERSION=2.1
+ARG PYTORCH_VERSION=2.8
 ARG MAX_ALLOWED_NCCL=2.22.3
+ARG MAKE_FLAGS="-j$(nproc)"
 
 SHELL ["/bin/bash", "-c"]
 # https://techoverflow.net/2019/05/18/how-to-fix-configuring-tzdata-interactive-input-when-building-docker-images/
@@ -30,8 +31,7 @@ ENV \
     PATH="$PATH:/root/.local/bin" \
     CUDA_TOOLKIT_ROOT_DIR="/usr/local/cuda" \
     MKL_THREADING_LAYER="GNU" \
-    # MAKEFLAGS="-j$(nproc)"
-    MAKEFLAGS="-j2"
+    MAKEFLAGS=${MAKE_FLAGS}
 
 RUN \
     CUDA_VERSION_MM=${CUDA_VERSION%.*} && \

dockers/nvidia/Dockerfile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG PYTORCH_VERSION=22.09
+ARG PYTORCH_VERSION=24.05
 
 # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes
 FROM nvcr.io/nvidia/pytorch:${PYTORCH_VERSION}-py3

dockers/release/Dockerfile

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 ARG PYTHON_VERSION=3.10
-ARG PYTORCH_VERSION=2.0
-ARG CUDA_VERSION=11.8.0
+ARG PYTORCH_VERSION=2.8
+ARG CUDA_VERSION=12.6.3
 
 FROM pytorchlightning/pytorch_lightning:base-cuda-py${PYTHON_VERSION}-torch${PYTORCH_VERSION}-cuda${CUDA_VERSION}