Allow torch 2.3 and remove torch ceiling version restriction (#2259)

* Allow torch 2.3 and remove torch ceiling version restriction

* Update base.py

* transformers quant test fix

---------

Co-authored-by: dbogunowicz <[email protected]>
Co-authored-by: Sara Adkins <[email protected]>

Author: 3 people

Diff for: docs/source/installation.md

@@ -18,7 +18,7 @@ limitations under the License.
 
 This repository is tested on Python 3.8-3.11, and Linux/Debian systems.
 It is recommended to install in a [virtual environment](https://docs.python.org/3/library/venv.html) to keep your system in order.
-Currently supported ML Frameworks are the following: `torch>=1.1.0,<1.14`, `tensorflow>=1.8.0,<=2.0.0`, `tensorflow.keras >= 2.2.0`.
+Currently supported ML Frameworks are the following: `torch>=1.7.0`, `tensorflow>=1.8.0,<=2.0.0`, `tensorflow.keras >= 2.2.0`.
 
 Install with pip using:
 

Diff for: setup.py

@@ -64,7 +64,7 @@
 
 _onnxruntime_deps = ["onnxruntime>=1.0.0"]
 _clip_deps = ["open_clip_torch==2.20.0"]
-supported_torch_version = "torch>=1.7.0,<2.3"
+supported_torch_version = "torch>=1.7.0"
 _pytorch_deps = [
     supported_torch_version,
     "gputils",

Diff for: src/sparseml/pytorch/base.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import functools
-import os
 from typing import Optional
 
 from sparseml.base import check_version
@@ -49,13 +48,11 @@
 
 
 _TORCH_MIN_VERSION = "1.0.0"
-# set max to 2.2.99 to account for bugfix versions with 2.2
-_TORCH_MAX_VERSION = os.environ.get("MAX_TORCH", "2.2.99")
 
 
 def check_torch_install(
     min_version: Optional[str] = _TORCH_MIN_VERSION,
-    max_version: Optional[str] = _TORCH_MAX_VERSION,
+    max_version: Optional[str] = None,
     raise_on_error: bool = True,
 ) -> bool:
     """
@@ -121,7 +118,7 @@ def check_torchvision_install(
 
 def require_torch(
     min_version: Optional[str] = _TORCH_MIN_VERSION,
-    max_version: Optional[str] = _TORCH_MAX_VERSION,
+    max_version: Optional[str] = None,
 ):
     """
     Decorator function to require use of torch.

Diff for: tests/sparseml/transformers/compression/test_quantization.py

@@ -53,7 +53,9 @@
 class TestQuantizationMatches(unittest.TestCase):
     old_recipe = None
     new_recipe = None
-    model_stub = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
+    # TODO: use "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T" for nightly
+    # or weekly runs, but this smaller model is better for commit testing
+    model_stub = "Xenova/llama2.c-stories15M"
     dataset = "open_platypus"
     old_output = "tiny_llama_old"
     new_output = "tiny_llama_new"
@@ -75,7 +77,7 @@ def setUpClass(cls):
         )
 
         cls.model_new = SparseAutoModelForCausalLM.from_pretrained(
-            cls.model_stub, device_map="cuda:1"
+            cls.model_stub, device_map="cuda:0"
         )
         cls._run_oneshot(
             cls.model_new,
@@ -106,6 +108,7 @@ def _run_oneshot(model, recipe, dataset, output_dir):
             num_calibration_samples=num_calibration_samples,
             recipe=recipe,
             pad_to_max_length=pad_to_max_length,
+            clear_sparse_session=True,
         )
 
     def _get_quant_info_old(self, model):
@@ -219,7 +222,7 @@ def test_perplexity(self):
         for idx, sample in enumerate(dataloader):
             if idx >= self.num_comparisons:
                 break
-            output_new = self.model_new(**tensors_to_device(sample, "cuda:1"))
+            output_new = self.model_new(**tensors_to_device(sample, "cuda:0"))
             output_old = self.model_old(**tensors_to_device(sample, "cuda:0"))
             if torch.isnan(output_old.loss) and torch.isnan(output_new.loss):
                 continue