CUDA 13.2 and torch 2.12.1 compatible Dockerfile. Volta (SM70) support dropped

CMakeLists.txt

@@ -8,7 +8,9 @@ endif()
 cmake_policy(SET CMP0057 NEW)
 
 # User-defined build options
-set(TORCHFORT_CUDA_CC_LIST "70;80;90" CACHE STRING "List of CUDA compute capabilities to build torchfort for.")
+# TORCHFORT_CUDA_CC_LIST default is set based on the detected CUDA Toolkit
+# version further below (after find_package(CUDAToolkit)). It can still be
+# overridden here on the command line via -DTORCHFORT_CUDA_CC_LIST=...
 set(TORCHFORT_NCCL_ROOT CACHE STRING "Path to search for NCCL installation. Default NVIDA HPC SDK provided NCCL version if available.")
 set(TORCHFORT_YAML_CPP_ROOT CACHE STRING "Path to search for yaml-cpp installation.")
 option(TORCHFORT_BUILD_FORTRAN "Build Fortran bindings" ON)
@@ -62,6 +64,22 @@ find_package(MPI REQUIRED)
 if (TORCHFORT_ENABLE_GPU)
   find_package(CUDAToolkit REQUIRED)
 
+  # Set default CUDA compute capabilities based on the detected CUDA Toolkit
+  # version if not provided by the user. Note: CUDA 13.0 dropped support for
+  # Volta (sm_70).
+  if (NOT TORCHFORT_CUDA_CC_LIST)
+    if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.0)
+      set(TORCHFORT_CUDA_CC_LIST_DEFAULTS "75;80;90;100")
+    elseif (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8)
+      set(TORCHFORT_CUDA_CC_LIST_DEFAULTS "70;80;90;100")
+    elseif (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 11.8)
+      set(TORCHFORT_CUDA_CC_LIST_DEFAULTS "70;80;90")
+    else()
+      set(TORCHFORT_CUDA_CC_LIST_DEFAULTS "70;80")
+    endif()
+    set(TORCHFORT_CUDA_CC_LIST ${TORCHFORT_CUDA_CC_LIST_DEFAULTS} CACHE STRING "List of CUDA compute capabilities to build torchfort for.")
+  endif()
+
   # HPC SDK
   # Locate and append NVHPC CMake configuration if available
   find_program(NVHPC_CXX_BIN "nvc++")

docker/Dockerfile

@@ -1,9 +1,9 @@
-FROM  nvcr.io/nvidia/nvhpc:25.7-devel-cuda12.9-ubuntu22.04
+FROM  nvcr.io/nvidia/nvhpc:26.5-devel-cuda13.2-ubuntu24.04
 
 SHELL ["/bin/bash", "-c"]
 
-ENV CUDA_HOME /opt/nvidia/hpc_sdk/Linux_x86_64/25.7/cuda
-ENV LD_LIBRARY_PATH /opt/nvidia/hpc_sdk/Linux_x86_64/25.7/cuda/12.9/extras/CUPTI/lib64:$LD_LIBRARY_PATH
+ENV CUDA_HOME /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda
+ENV LD_LIBRARY_PATH /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/cuda/13.2/extras/CUPTI/lib64:$LD_LIBRARY_PATH
 
 # Install System Dependencies
 ENV DEBIAN_FRONTEND noninteractive
@@ -13,14 +13,15 @@ RUN apt update -y && \
     apt install -y vim gfortran
 
 # Install PyTorch (and select dependencies)
-RUN pip3 install --no-deps torch==2.8.0 --index-url https://download.pytorch.org/whl/cu129 && \
-    pip3 install --no-deps nvidia-cudnn-cu12==9.10.2.21 nvidia-cusparselt-cu12==0.7.1 nvidia-cufile-cu12==1.14.1.1 nvidia-nccl-cu12==2.27.3
+RUN pip3 install --break-system-packages typing_extensions && \
+    pip3 install --break-system-packages --no-deps torch==2.12.1 --index-url https://download.pytorch.org/whl/cu132 && \
+    pip3 install --break-system-packages --no-deps nvidia-cudnn-cu13==9.20.0.48 nvidia-cusparselt-cu13==0.8.1 nvidia-nccl-cu13==2.29.7 nvidia-cufile==1.18.1.6
 
 # Remove conflicting NCCL version from NVHPC SDK, add libnccl.so symlink to pip installed NCCL
-RUN rm -rf /opt/nvidia/hpc_sdk/Linux_x86_64/25.7/comm_libs/nccl/ && \
-    cd /usr/local/lib/python3.10/dist-packages/nvidia/nccl/lib && \
+RUN rm -rf /opt/nvidia/hpc_sdk/Linux_x86_64/26.5/comm_libs/nccl/ && \
+    cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
     ln -s libnccl.so.2 libnccl.so
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.10/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
 
 # Install yaml-cpp
 RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
@@ -34,7 +35,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
 ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
 
 # Install additional Python dependencies
-RUN pip3 install wandb ruamel-yaml matplotlib pygame moviepy
+RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
 
 # Install TorchFort
 ENV FC=nvfortran
@@ -44,12 +45,12 @@ RUN cd /torchfort && mkdir build && cd build && \
     cmake -DCMAKE_INSTALL_PREFIX=/opt/torchfort \
     -DCMAKE_CXX_COMPILER=`which g++` \
     -DTORCHFORT_YAML_CPP_ROOT=/opt/yaml-cpp \
-    -DTORCHFORT_NCCL_ROOT=/usr/local/lib/python3.10/dist-packages/nvidia/nccl \
+    -DTORCHFORT_NCCL_ROOT=/usr/local/lib/python3.12/dist-packages/nvidia/nccl \
     -DTORCHFORT_BUILD_EXAMPLES=1 \
     -DTORCHFORT_BUILD_TESTS=1 \
     -DCMAKE_PREFIX_PATH="`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`" \
     .. && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
 ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.10/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}

docker/Dockerfile_gnu

@@ -1,4 +1,4 @@
-FROM nvcr.io/nvidia/cuda:12.9.1-devel-ubuntu22.04
+FROM nvcr.io/nvidia/cuda:13.2.0-devel-ubuntu24.04
 
 SHELL ["/bin/bash", "-c"]
 
@@ -14,19 +14,27 @@ RUN apt update -y && \
 
 # Download HPCX
 RUN cd /opt && \
-    wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda12/hpcx-v2.24.1-gcc-doca_ofed-ubuntu22.04-cuda12-x86_64.tbz && \
-    tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu22.04-cuda12-x86_64.tbz && \
-    mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu22.04-cuda12-x86_64 hpcx && \
-    cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu22.04-cuda12-x86_64.tbz
+    wget https://content.mellanox.com/hpc/hpc-x/v2.24.1_cuda13/hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
+    tar xjf hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz && \
+    mv hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64 hpcx && \
+    cd /opt && rm hpcx-v2.24.1-gcc-doca_ofed-ubuntu24.04-cuda13-x86_64.tbz
 
 ENV PATH /opt/hpcx/ompi/bin:$PATH
 ENV LD_LIBRARY_PATH /opt/hpcx/ompi/lib:$LD_LIBRARY_PATH
 
 RUN echo "source /opt/hpcx/hpcx-init.sh; hpcx_load" >>  /root/.bashrc
 
-# Install PyTorch
-RUN pip3 install --no-deps torch==2.8.0 --index-url https://download.pytorch.org/whl/cu129 && \
-    pip3 install --no-deps nvidia-cudnn-cu12==9.10.2.21 nvidia-cusparselt-cu12==0.7.1 nvidia-cufile-cu12==1.14.1.1
+# Install PyTorch (and select dependencies)
+RUN pip3 install --break-system-packages typing_extensions && \
+    pip3 install --break-system-packages --no-deps torch==2.12.1 --index-url https://download.pytorch.org/whl/cu132 && \
+    pip3 install --break-system-packages --no-deps nvidia-cudnn-cu13==9.20.0.48 nvidia-cusparselt-cu13==0.8.1 nvidia-nccl-cu13==2.29.7 nvidia-cufile==1.18.1.6 nvidia-nvshmem-cu13==3.4.5
+
+# Add libnccl.so symlink to pip installed NCCL (CUDA devel image ships no NCCL)
+RUN cd /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib && \
+    ln -s libnccl.so.2 libnccl.so
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib:$LD_LIBRARY_PATH
+# NVSHMEM is a torch runtime dependency (libnvshmem_host.so.3) not present in the CUDA image
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/nvidia/nvshmem/lib:$LD_LIBRARY_PATH
 
 # Install yaml-cpp
 RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
@@ -40,7 +48,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
 ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
 
 # Install additional Python dependencies
-RUN pip3 install wandb ruamel-yaml matplotlib pygame moviepy
+RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
 
 # Install TorchFort
 ENV FC=gfortran
@@ -49,11 +57,12 @@ RUN source /opt/hpcx/hpcx-init.sh &&  hpcx_load && \
     cd /torchfort && mkdir build && cd build && \
     cmake -DCMAKE_INSTALL_PREFIX=/opt/torchfort \
     -DTORCHFORT_YAML_CPP_ROOT=/opt/yaml-cpp \
+    -DTORCHFORT_NCCL_ROOT=/usr/local/lib/python3.12/dist-packages/nvidia/nccl \
     -DTORCHFORT_BUILD_EXAMPLES=1 \
     -DTORCHFORT_BUILD_TESTS=1 \
     -DCMAKE_PREFIX_PATH="`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`" \
     .. && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
 ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.10/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}

docker/Dockerfile_gnu_cpuonly

@@ -1,4 +1,4 @@
-FROM ubuntu:22.04
+FROM ubuntu:24.04
 
 SHELL ["/bin/bash", "-c"]
 
@@ -13,7 +13,7 @@ RUN apt update -y && \
     apt install -y openmpi-bin libopenmpi-dev
 
 # Install PyTorch
-RUN pip3 install torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu
+RUN pip3 install --break-system-packages torch==2.12.1 --index-url https://download.pytorch.org/whl/cpu
 
 # Install yaml-cpp
 RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
@@ -27,7 +27,7 @@ RUN git clone https://github.com/jbeder/yaml-cpp.git --branch 0.8.0 && \
 ENV LD_LIBRARY_PATH /opt/yaml-cpp/lib:${LD_LIBRARY_PATH}
 
 # Install additional Python dependencies
-RUN pip3 install wandb ruamel-yaml matplotlib pygame moviepy
+RUN pip3 install --break-system-packages wandb ruamel-yaml matplotlib pygame moviepy
 
 # Install TorchFort without GPU support
 ENV FC=gfortran
@@ -43,4 +43,4 @@ RUN cd /torchfort && mkdir build && cd build && \
     make -j$(nproc) install && \
     cd / && rm -rf torchfort
 ENV LD_LIBRARY_PATH /opt/torchfort/lib:${LD_LIBRARY_PATH}
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.10/dist-packages/torch/lib:${LD_LIBRARY_PATH}
+ENV LD_LIBRARY_PATH /usr/local/lib/python3.12/dist-packages/torch/lib:${LD_LIBRARY_PATH}

requirements.txt

@@ -2,7 +2,7 @@
 ruamel-yaml
 
 # pytorch and some dependencies
-torch==2.8.0
+torch==2.12.1
 
 # training monitoring
 wandb

src/csrc/cuda_wrap.cpp

@@ -25,8 +25,8 @@
 #define LOAD_SYM(symbol, version, optional)                                                                            \
   do {                                                                                                                 \
     cudaDriverEntryPointQueryResult driverStatus = cudaDriverEntryPointSymbolNotFound;                                 \
-    cudaError_t err = cudaGetDriverEntryPointByVersion(#symbol, (void**)(&cuFnTable.pfn_##symbol), version,   \
-                                                       cudaEnableDefault, &driverStatus));                             \
+    cudaError_t err = cudaGetDriverEntryPointByVersion(#symbol, (void**)(&cuFnTable.pfn_##symbol), version,            \
+                                                       cudaEnableDefault, &driverStatus);                              \
     if ((driverStatus != cudaDriverEntryPointSuccess || err != cudaSuccess) && !optional) {                            \
       THROW_CUDA_ERROR("cudaGetDriverEntryPointByVersion failed.");                                                    \
     }                                                                                                                  \