Merge remote-tracking branch 'origin/main' into feature/damian/kv_cac…

…he_codegen

.MAINTAINERS

@@ -9,10 +9,9 @@ rahul-tuli
 KSGulin
 dbogunowicz
 # anmarques
-kylesayrs
 # eldarkurtic
-corey-nm
 # chibukach
 # shubhra
 # abhinavnmagic
 # eiofinov
+dsikka

.github/workflows/test-check.yaml

@@ -16,10 +16,8 @@ jobs:
       branch: ${{ steps.get-branch.outputs.branch }}
       base: ${{ steps.base-check.outputs.output }}
       deepsparse: ${{ steps.deepsparse-check.outputs.output }}
-      keras: ${{ steps.keras-check.outputs.output }}
       onnx: ${{ steps.onnx-check.outputs.output }}
       pytorch: ${{ steps.pytorch-check.outputs.output }}
-      tensorflow_v1: ${{ steps.tensorflow_v1-check.outputs.output }}
     steps:
       - uses: actions/checkout@v2
         with:
@@ -43,12 +41,6 @@ jobs:
           ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/deepsparse|setup.py|.github")
           || (echo $GITHUB_REF | grep -E "refs/heads/[release/|main]"))
           && echo "::set-output name=output::1" || echo "::set-output name=output::0"
-      - name: "Checking if sparseml.keras was changed"
-        id: keras-check
-        run: >
-          ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/keras|setup.py|.github")
-          || (echo $GITHUB_REF | grep -E "refs/heads/[release/|main]"))
-          && echo "::set-output name=output::1" || echo "::set-output name=output::0"
       - name: "Checking if sparseml.onnx was changed"
         id: onnx-check
         run: >
@@ -61,12 +53,6 @@ jobs:
           ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/pytorch|setup.py|.github")
           || (echo $GITHUB_REF | grep -E "refs/heads/[release/|main]"))
           && echo "::set-output name=output::1" || echo "::set-output name=output::0"
-      - name: "Checking if sparseml.tensorflow_v1 was changed"
-        id: tensorflow_v1-check
-        run: >
-          ((git diff --name-only origin/main HEAD | grep -E "[src|tests]/sparseml/tensorflow_v1|setup.py|.github")
-          || (echo $GITHUB_REF | grep -E "refs/heads/[release/|main]"))
-          && echo "::set-output name=output::1" || echo "::set-output name=output::0"
   base-tests:
     runs-on: ubuntu-22.04
     env:
@@ -119,30 +105,6 @@ jobs:
         run: pip3 install .[dev,deepsparse,onnxruntime]
       - name: "🔬 Running deepsparse tests"
         run: make test TARGETS=deepsparse
-  keras-tests:
-    runs-on: ubuntu-22.04
-    env:
-      SPARSEZOO_TEST_MODE: "true"
-    needs: test-setup
-    if:  ${{needs.test-setup.outputs.keras == 1}}
-    steps:
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.8'
-      - uses: actions/checkout@v2
-      - uses: actions/checkout@v2
-        with:
-          repository: "neuralmagic/sparsezoo"
-          path: "sparsezoo"
-          ref: ${{needs.test-setup.outputs.branch}}
-      - name: "⚙️ Install sparsezoo dependencies"
-        run: pip3 install -U pip && pip3 install setuptools sparsezoo/
-      - name: "Clean sparsezoo directory"
-        run: rm -r sparsezoo/
-      - name: "⚙️ Install dependencies"
-        run: pip3 install .[dev,tf_keras,onnxruntime]
-      - name: "🔬 Running keras tests"
-        run: make test TARGETS=keras
   onnx-tests:
     runs-on: ubuntu-22.04
     env:
@@ -191,30 +153,6 @@ jobs:
         run: pip3 install .[dev,torchvision,onnxruntime]
       - name: "🔬 Running pytorch tests"
         run: make test TARGETS=pytorch
-  tensorflow-v1-tests:
-    runs-on: ubuntu-22.04
-    env:
-      SPARSEZOO_TEST_MODE: "true"
-    needs: test-setup
-    if:  ${{needs.test-setup.outputs.tensorflow_v1 == 1}}
-    steps:
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.7'
-      - uses: actions/checkout@v2
-      - uses: actions/checkout@v2
-        with:
-          repository: "neuralmagic/sparsezoo"
-          path: "sparsezoo"
-          ref: ${{needs.test-setup.outputs.branch}}
-      - name: "⚙️ Install sparsezoo dependencies"
-        run: pip3 install -U pip && pip3 install setuptools sparsezoo/
-      - name: "Clean sparsezoo directory"
-        run: rm -r sparsezoo/
-      - name: "⚙️ Install dependencies"
-        run: pip3 install .[dev,tf_v1,onnxruntime]
-      - name: "🔬 Running tensorflow_v1 tests"
-        run: make test TARGETS=tensorflow_v1
   compat-pytorch-1_9-pytorch-tests:
     runs-on: ubuntu-22.04
     env:

DEVELOPING.md

@@ -16,7 +16,7 @@ limitations under the License.
 
 # Developing SparseML
 
-SparseML is developed and tested using Python 3.7-3.10.
+SparseML is developed and tested using Python 3.8-3.10.
 To develop SparseML, you will also need the development dependencies and to follow the styling guidelines.
 
 Here are some details to get started.
@@ -40,6 +40,9 @@ Those can be installed by adding the framework name to the install extras. Frame
 python3 -m pip install -e "./[dev,torch]"
 ```
 
+Note: Running all pytorch tests using `make test TARGETS=torch`, also requires `torchvision`
+and `onnxruntime` install all these dependencies using `python3 -m pip install -e "./[dev, torch, torchvision, onnxruntime]"`
+
 **Code Styling and Formatting checks**
 
 ```bash

README.md

@@ -104,10 +104,10 @@ SparseML enables you to create a sparse model trained on your dataset in two way
 
 ## Installation
 
-This repository is tested on Python 3.7-3.10, and Linux/Debian systems.
+This repository is tested on Python 3.8-3.10, 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.12.1`, `tensorflow>=1.8.0,<2.0.0`, `tensorflow.keras >= 2.2.0`.
+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`.
 
 Install with pip using:
 

docs/source/installation.md

@@ -16,9 +16,9 @@ limitations under the License.
 
 # Installation
 
-This repository is tested on Python 3.7-3.10, and Linux/Debian systems.
+This repository is tested on Python 3.8-3.10, 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.8.0`, `tensorflow>=1.8.0,<=2.0.0`, `tensorflow.keras >= 2.2.0`.
+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`.
 
 Install with pip using:
 

integrations/ultralytics-yolov5/tutorials/sparse-transfer-learning.md

@@ -31,7 +31,7 @@ SparseZoo contains pre-sparsified checkpoints of each YOLOv5 model. These models
 Install via `pip`:
 
 ```
-pip install sparseml[torchvision]
+pip install "sparseml[yolov5]"
 ```
 
 ## Table of Contents

research/information_retrieval/doc2query/requirements.txt

@@ -94,7 +94,7 @@ pyzmq==22.1.0
 qtconsole==5.1.0
 QtPy==1.9.0
 regex==2021.4.4
-requests==2.25.1
+requests==2.31.0
 rouge-score==0.0.4
 sacremoses==0.0.45
 scikit-image==0.18.1
@@ -121,7 +121,7 @@ tifffile==2021.6.14
 tokenizers==0.10.3
 toposort==1.6
 torch==1.8.0
-tornado==6.1
+tornado==6.3.2
 tqdm==4.61.1
 traitlets==5.0.5
 transformers==4.7.0

setup.py

@@ -47,7 +47,8 @@
     "requests>=2.0.0",
     "scikit-image>=0.15.0",
     "scikit-learn>=0.24.2",
-    "scipy>=1.0.0",
+    "scipy<1.9.2,>=1.8; python_version <= '3.9'",
+    "scipy>=1.0.0; python_version > '3.9'",
     "tqdm>=4.0.0",
     "toposort>=1.0",
     "GPUtil>=1.4.0",
@@ -61,16 +62,17 @@
 _deepsparse_ent_deps = [f"deepsparse-ent~={version_nm_deps}"]
 
 _onnxruntime_deps = ["onnxruntime>=1.0.0"]
+supported_torch_version = "torch>=1.7.0,<=2.0"
 _pytorch_deps = [
-    "torch>=1.1.0,<=1.13.1",
+    supported_torch_version,
     "gputils",
 ]
 _pytorch_all_deps = _pytorch_deps + [
-    "torchvision>=0.3.0,<0.15",
-    "torchaudio<=0.13",
+    "torchvision>=0.3.0,<=0.15.1",
+    "torchaudio<=2.0.1",
 ]
 _pytorch_vision_deps = _pytorch_deps + [
-    "torchvision>=0.3.0,<0.15",
+    "torchvision>=0.3.0,<=0.15.1",
     "opencv-python<=4.6.0.66",
 ]
 _transformers_deps = _pytorch_deps + [
@@ -118,7 +120,11 @@
     "tensorboardX>=1.0",
 ]
 
-_ultralytics_deps = ["ultralytics==8.0.30"]
+
+_ultralytics_deps = [
+    "ultralytics==8.0.30",
+    supported_torch_version,
+]
 
 
 def _setup_packages() -> List:
@@ -284,7 +290,7 @@ def _setup_long_description() -> Tuple[str, str]:
     install_requires=_setup_install_requires(),
     extras_require=_setup_extras(),
     entry_points=_setup_entry_points(),
-    python_requires=">=3.7.0,<3.11",
+    python_requires=">=3.8.0,<3.11",
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Programming Language :: Python :: 3",

src/sparseml/exporters/kv_cache_injector.py

@@ -31,6 +31,7 @@ def __init__(
         self,
         model_path: str,
         inplace: bool = True,
+        **kwargs: Any,
     ):
         """
         A transformation that injects Key Value cache support into the model.
@@ -45,11 +46,25 @@ def __init__(
         This transformation not only injects the cache support, but also adjusts
         the model to account for the cache support. This means altering the input
         to the model, such as adding "position" input to the model.
+        to determine the values:
+        - num_attention_heads
+        - hidden_size_kv_cache
+        required to enforce static dimensions of the kv cache input/output.
+        If `model_path` is not provided, the two aforementioned values must
+        be provided in the `kwargs`.
+
+        This transformation not only injects the cache support, but also adjusts
+        the model to account for the cache support. This means altering the input
+        to the model, such as altering the positions to account for the injected
+        key/value pairs.
 
         Usage:
         ```python
         onnx_model: onnx.ModelProto = ...
         exporter = KeyValueCacheInjector(model_path="path/to/model")
+        # alternatively
+        # exporter = KeyValueCacheInjector(num_attention_heads = 16,
+        #                                  hidden_size_dim = 64)
         exporter.export(onnx_model, "model.onnx")
         ```
 
@@ -67,7 +82,7 @@ def __init__(
         self.inplace = inplace
         self.config = get_kv_cache_config(model_path)
         PositionEmbeddingAdjustment = self.config.positional_embedding_transform
-
+        # add kwargs support
         transforms = [
             CacheKeysAndValues(
                 num_attention_heads=self.config.num_attention_heads,
@@ -77,6 +92,7 @@ def __init__(
                 transpose_kv_key_input=self.config.transpose_kv_key_input,
             ),
             PositionEmbeddingAdjustment(),
+
         ]
         super().__init__(transforms)
 

src/sparseml/exporters/transforms/kv_cache/cache_keys_and_values.py

@@ -141,9 +141,9 @@ class CacheKeysAndValues(OnnxTransform):
 
     def __init__(
         self,
-        num_attention_heads: int = 16,
-        hidden_size_kv_cache: int = 64,
-        internally_multiply_batch_by_num_attention_heads: bool = True,
+        num_attention_heads: int,
+        hidden_size_kv_cache: int,
+        internally_multiply_batch_by_num_attention_heads: bool,
         transpose_kv_value_input: Optional[Tuple[int, int, int, int]] = None,
         transpose_kv_key_input: Optional[Tuple[int, int, int, int]] = None,
     ):