upgrade libtorch to 2.6.0 and cutlass to 3.8.0 (#398)

Author: guocuimi

.github/workflows/build_wheel.yml

@@ -22,12 +22,9 @@ jobs:
   build_wheel:
     strategy:
       matrix:
-        python: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python: ["3.9", "3.10", "3.11", "3.12"]
         cuda: ["11.8", "12.1", "12.4"]
-        torch: ["2.4.1", "2.5.1"]
-        exclude: # torch 2.5.1 dropped support for python 3.8
-          - python: "3.8"
-            torch: "2.5.1"
+        torch: ["2.4.1", "2.5.1", "2.6.0"]
     runs-on: [self-hosted, linux, release]
     env:
       PYTHON_VERSION: ${{ matrix.python }}

.github/workflows/package_test.yml

@@ -40,7 +40,7 @@ jobs:
       matrix:
         python: ["3.12"]
         cuda: ["12.4"]
-        torch: ["2.5.1"]
+        torch: ["2.6.0"]
     runs-on: [self-hosted, linux, build]
     env:
       PYTHON_VERSION: ${{ matrix.python }}

.github/workflows/publish_wheel.yml

@@ -23,7 +23,7 @@ jobs:
       matrix:
         python: ["3.9", "3.10", "3.11", "3.12"]
         cuda: ["12.4"]
-        torch: ["2.5.1"]
+        torch: ["2.6.0"]
     runs-on: [self-hosted, linux, release]
     env:
       PYTHON_VERSION: ${{ matrix.python }}

.github/workflows/release_test.yml

@@ -21,7 +21,7 @@ jobs:
       matrix:
         python: ["3.9", "3.10", "3.11", "3.12"]
         cuda: ["12.4"]
-        torch: ["2.5.1"]
+        torch: ["2.6.0"]
     runs-on: [self-hosted, linux, release]
     env:
       PYTHON_VERSION: ${{ matrix.python }}

CMakeLists.txt

@@ -196,25 +196,25 @@ if (DEFINED ENV{LIBTORCH_ROOT})
 else()
   include(FetchContent)
   if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.4)
-    # download libtorch 2.5.1 with cuda 12.4 from pytorch.org
+    # download libtorch 2.6.0 with cuda 12.4 from pytorch.org
     if (USE_CXX11_ABI)
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcu124.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu124.zip")
     else()
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-shared-with-deps-2.5.1%2Bcu124.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu124/libtorch-shared-with-deps-2.6.0%2Bcu124.zip")
     endif()
   elseif(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.1)
-    # download libtorch 2.5.1 with cuda 12.1 from pytorch.org
+    # download libtorch 2.6.0 with cuda 12.1 from pytorch.org
     if (USE_CXX11_ABI)
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcu121.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu121.zip")
     else()
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.5.1%2Bcu121.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.6.0%2Bcu121.zip")
     endif()
   elseif(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 11.8)
-    # download libtorch 2.5.1 with cuda 11.8 from pytorch.org
+    # download libtorch 2.6.0 with cuda 11.8 from pytorch.org
     if (USE_CXX11_ABI)
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.5.1%2Bcu118.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.6.0%2Bcu118.zip")
     else()
-      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu118/libtorch-shared-with-deps-2.5.1%2Bcu118.zip")
+      set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cu118/libtorch-shared-with-deps-2.6.0%2Bcu118.zip")
     endif()
   else()
     # error out if cuda version is not supported
@@ -234,7 +234,7 @@ else()
   FetchContent_MakeAvailable(libtorch)
 
   find_package(Torch REQUIRED PATHS ${libtorch_SOURCE_DIR} NO_DEFAULT_PATH)
-  message(STATUS "Downloading and using libtorch 2.5.1 for cuda ${CUDA_VERSION} at ${libtorch_SOURCE_DIR}")
+  message(STATUS "Downloading and using libtorch 2.6.0 for cuda ${CUDA_VERSION} at ${libtorch_SOURCE_DIR}")
 endif()
 
 # check if USE_CXX11_ABI is set correctly

README.md

@@ -55,7 +55,7 @@ ScaleLLM is currently undergoing active development. We are fully committed to c
 
 ScaleLLM is available as a Python Wheel package on PyPI. You can install it using pip:
 ```bash
-# Install scalellm with CUDA 12.4 and Pytorch 2.5.1
+# Install scalellm with CUDA 12.4 and Pytorch 2.6.0
 pip install -U scalellm
 ```
 

docs/source/index.rst

@@ -12,7 +12,7 @@ ScaleLLM is available as a Python Wheel package on `PyPI <https://pypi.org/proje
 
 .. code-block:: bash
 
-   # Install ScaleLLM with CUDA 12.4 and Pytorch 2.5.1
+   # Install ScaleLLM with CUDA 12.4 and Pytorch 2.6.0
    $ pip install -U scalellm
 
 

docs/source/quick_start.rst

@@ -25,6 +25,12 @@ If you want to install ScaleLLM with different versions of CUDA and PyTorch, you
 
         .. tabs::
 
+            .. tab:: PyTorch 2.6.0
+
+                .. code-block:: bash
+
+                    $ pip install -U scalellm -i https://whl.vectorch.com/cu124/torch2.6.0/
+
             .. tab:: PyTorch 2.5.1
 
                 .. code-block:: bash
@@ -41,6 +47,12 @@ If you want to install ScaleLLM with different versions of CUDA and PyTorch, you
 
         .. tabs::
 
+            .. tab:: PyTorch 2.6.0
+
+                .. code-block:: bash
+
+                    $ pip install -U scalellm -i https://whl.vectorch.com/cu121/torch2.6.0/
+
             .. tab:: PyTorch 2.5.1
 
                 .. code-block:: bash
@@ -57,6 +69,12 @@ If you want to install ScaleLLM with different versions of CUDA and PyTorch, you
 
         .. tabs::
 
+            .. tab:: PyTorch 2.6.0
+
+                .. code-block:: bash
+
+                    $ pip install -U scalellm -i https://whl.vectorch.com/cu118/torch2.6.0/
+
             .. tab:: PyTorch 2.5.1
 
                 .. code-block:: bash

scalellm/downloader.py

@@ -16,7 +16,12 @@ def convert_pickle_to_safetensors(path):
                 continue
 
             # load the model
-            model = torch.load(file_path, map_location="cpu")
+            try:
+                model = torch.load(file_path, map_location="cpu")
+            except Exception as e:
+                print(f"Error loading {filename}: {e}")
+                continue
+
             if hasattr(model, "state_dict"):
                 state_dict = model.state_dict()
             else: