From 10a8f2d41fc7a6af2e2f3ac53d05d7c8d95a8c87 Mon Sep 17 00:00:00 2001 From: dbogunowicz Date: Wed, 17 Apr 2024 13:13:20 +0000 Subject: [PATCH] initial commit --- README.md | 2 +- setup.py | 4 ++-- src/sparsetensors/README.md | 2 +- src/sparsetensors/compressors/base.py | 6 +++--- src/sparsetensors/compressors/dense.py | 2 +- src/sparsetensors/compressors/sparse_bitmask.py | 4 ++-- src/sparsetensors/config/base.py | 2 +- src/sparsetensors/config/dense.py | 2 +- src/sparsetensors/config/sparse_bitmask.py | 2 +- src/sparsetensors/quantization/lifecycle/apply.py | 10 +++++----- .../quantization/lifecycle/calibration.py | 2 +- src/sparsetensors/quantization/lifecycle/forward.py | 6 +++--- src/sparsetensors/quantization/lifecycle/frozen.py | 2 +- src/sparsetensors/quantization/lifecycle/initialize.py | 8 ++++---- src/sparsetensors/quantization/observers/base.py | 4 ++-- src/sparsetensors/quantization/observers/helpers.py | 2 +- src/sparsetensors/quantization/observers/memoryless.py | 4 ++-- src/sparsetensors/quantization/observers/min_max.py | 6 +++--- src/sparsetensors/quantization/quant_args.py | 2 +- src/sparsetensors/quantization/quant_config.py | 4 ++-- src/sparsetensors/quantization/quant_scheme.py | 2 +- src/sparsetensors/utils/helpers.py | 6 +++--- tests/quantization/lifecycle/test_apply.py | 4 ++-- tests/quantization/test_quant_args.py | 2 +- tests/quantization/test_quant_config.py | 2 +- tests/quantization/test_quant_scheme.py | 2 +- tests/test_bitmask.py | 2 +- tests/test_registry.py | 2 +- 28 files changed, 49 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index 53d7e9bd..05fa83a3 100644 --- a/README.md +++ b/README.md @@ -1 +1 @@ -# sparsetensors +# compressed-tensors diff --git a/setup.py b/setup.py index 303f14ad..6ab8dc09 100644 --- a/setup.py +++ b/setup.py @@ -18,7 +18,7 @@ def _setup_packages() -> List: return find_packages( - "src", include=["sparsetensors", "sparsetensors.*"], exclude=["*.__pycache__.*"] + "src", include=["compressed-tensors", "compressed-tensors.*"], exclude=["*.__pycache__.*"] ) def _setup_install_requires() -> List: @@ -28,7 +28,7 @@ def _setup_extras() -> Dict: return {"dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0",]} setup( - name="sparsetensors", + name="compressed-tensors", version="0.3.0", author="Neuralmagic, Inc.", author_email="support@neuralmagic.com", diff --git a/src/sparsetensors/README.md b/src/sparsetensors/README.md index 833d1ec2..5b1c8ece 100644 --- a/src/sparsetensors/README.md +++ b/src/sparsetensors/README.md @@ -65,7 +65,7 @@ is saved to the model's `config.json` file. **Note:** the model must have been initialized with SparseAutoModelForCausalLM.from_pretrained() ```python -from sparsetensors import BitmaskConfig +from compressed_tensors import BitmaskConfig output_dir = "/PATH/TO/SAVE/COMPRESSED_MODEL" sparsity_config = BitmaskConfig() diff --git a/src/sparsetensors/compressors/base.py b/src/sparsetensors/compressors/base.py index a82ce048..9c205f93 100644 --- a/src/sparsetensors/compressors/base.py +++ b/src/sparsetensors/compressors/base.py @@ -15,9 +15,9 @@ import operator from typing import Dict, Generator, Tuple -from sparsetensors.base import SPARSITY_CONFIG_NAME -from sparsetensors.config import CompressionConfig -from sparsetensors.registry import RegistryMixin +from compressed_tensors.base import SPARSITY_CONFIG_NAME +from compressed_tensors.config import CompressionConfig +from compressed_tensors.registry import RegistryMixin from torch import Tensor from torch.nn import Module, Parameter from tqdm import tqdm diff --git a/src/sparsetensors/compressors/dense.py b/src/sparsetensors/compressors/dense.py index 87f112ac..6e8785bc 100644 --- a/src/sparsetensors/compressors/dense.py +++ b/src/sparsetensors/compressors/dense.py @@ -14,7 +14,7 @@ from typing import Dict, Generator, Tuple -from sparsetensors.compressors import ModelCompressor +from compressed_tensors.compressors import ModelCompressor from torch import Tensor diff --git a/src/sparsetensors/compressors/sparse_bitmask.py b/src/sparsetensors/compressors/sparse_bitmask.py index 3043e43b..9edf0b69 100644 --- a/src/sparsetensors/compressors/sparse_bitmask.py +++ b/src/sparsetensors/compressors/sparse_bitmask.py @@ -18,8 +18,8 @@ import numpy import torch from safetensors import safe_open -from sparsetensors.compressors import ModelCompressor -from sparsetensors.utils import get_nested_weight_mappings, merge_names +from compressed_tensors.compressors import ModelCompressor +from compressed_tensors.utils import get_nested_weight_mappings, merge_names from torch import Tensor from tqdm import tqdm diff --git a/src/sparsetensors/config/base.py b/src/sparsetensors/config/base.py index 8af48bd9..552c10a1 100644 --- a/src/sparsetensors/config/base.py +++ b/src/sparsetensors/config/base.py @@ -15,7 +15,7 @@ from typing import Optional from pydantic import BaseModel -from sparsetensors.registry import RegistryMixin +from compressed_tensors.registry import RegistryMixin __all__ = ["CompressionConfig"] diff --git a/src/sparsetensors/config/dense.py b/src/sparsetensors/config/dense.py index a37be308..aa23220c 100644 --- a/src/sparsetensors/config/dense.py +++ b/src/sparsetensors/config/dense.py @@ -14,7 +14,7 @@ from typing import Optional -from sparsetensors.config import CompressionConfig +from compressed_tensors.config import CompressionConfig __all__ = ["DenseSparsityConfig"] diff --git a/src/sparsetensors/config/sparse_bitmask.py b/src/sparsetensors/config/sparse_bitmask.py index d17c6a1a..9b9cf211 100644 --- a/src/sparsetensors/config/sparse_bitmask.py +++ b/src/sparsetensors/config/sparse_bitmask.py @@ -14,7 +14,7 @@ from typing import Optional -from sparsetensors.config.base import CompressionConfig +from compressed_tensors.config.base import CompressionConfig __all__ = ["BitmaskConfig"] diff --git a/src/sparsetensors/quantization/lifecycle/apply.py b/src/sparsetensors/quantization/lifecycle/apply.py index ac238564..77c5245b 100644 --- a/src/sparsetensors/quantization/lifecycle/apply.py +++ b/src/sparsetensors/quantization/lifecycle/apply.py @@ -16,16 +16,16 @@ from collections import OrderedDict from typing import Iterable, Optional -from sparsetensors.quantization.lifecycle.calibration import set_module_for_calibration -from sparsetensors.quantization.lifecycle.frozen import freeze_module_quantization -from sparsetensors.quantization.lifecycle.initialize import ( +from compressed_tensors.quantization.lifecycle.calibration import set_module_for_calibration +from compressed_tensors.quantization.lifecycle.frozen import freeze_module_quantization +from compressed_tensors.quantization.lifecycle.initialize import ( initialize_module_for_quantization, ) -from sparsetensors.quantization.quant_config import ( +from compressed_tensors.quantization.quant_config import ( QuantizationConfig, QuantizationStatus, ) -from sparsetensors.quantization.utils import iter_named_leaf_modules +from compressed_tensors.quantization.utils import iter_named_leaf_modules from torch.nn import Module diff --git a/src/sparsetensors/quantization/lifecycle/calibration.py b/src/sparsetensors/quantization/lifecycle/calibration.py index 51c594fb..7ab1d896 100644 --- a/src/sparsetensors/quantization/lifecycle/calibration.py +++ b/src/sparsetensors/quantization/lifecycle/calibration.py @@ -15,7 +15,7 @@ import logging -from sparsetensors.quantization.quant_config import QuantizationStatus +from compressed_tensors.quantization.quant_config import QuantizationStatus from torch.nn import Module diff --git a/src/sparsetensors/quantization/lifecycle/forward.py b/src/sparsetensors/quantization/lifecycle/forward.py index 6416a10b..e9db69ae 100644 --- a/src/sparsetensors/quantization/lifecycle/forward.py +++ b/src/sparsetensors/quantization/lifecycle/forward.py @@ -15,9 +15,9 @@ from functools import wraps import torch -from sparsetensors.quantization.quant_args import QuantizationArgs -from sparsetensors.quantization.quant_config import QuantizationStatus -from sparsetensors.quantization.quant_scheme import QuantizationScheme +from compressed_tensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.quantization.quant_config import QuantizationStatus +from compressed_tensors.quantization.quant_scheme import QuantizationScheme from torch.nn import Module diff --git a/src/sparsetensors/quantization/lifecycle/frozen.py b/src/sparsetensors/quantization/lifecycle/frozen.py index 63949cf5..9df8ec46 100644 --- a/src/sparsetensors/quantization/lifecycle/frozen.py +++ b/src/sparsetensors/quantization/lifecycle/frozen.py @@ -13,7 +13,7 @@ # limitations under the License. -from sparsetensors.quantization.quant_config import QuantizationStatus +from compressed_tensors.quantization.quant_config import QuantizationStatus from torch.nn import Module diff --git a/src/sparsetensors/quantization/lifecycle/initialize.py b/src/sparsetensors/quantization/lifecycle/initialize.py index aa6e3994..d78997c1 100644 --- a/src/sparsetensors/quantization/lifecycle/initialize.py +++ b/src/sparsetensors/quantization/lifecycle/initialize.py @@ -17,10 +17,10 @@ from typing import Optional import torch -from sparsetensors.quantization.lifecycle.forward import wrap_module_forward_quantized -from sparsetensors.quantization.quant_args import QuantizationArgs -from sparsetensors.quantization.quant_config import QuantizationStatus -from sparsetensors.quantization.quant_scheme import QuantizationScheme +from compressed_tensors.quantization.lifecycle.forward import wrap_module_forward_quantized +from compressed_tensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.quantization.quant_config import QuantizationStatus +from compressed_tensors.quantization.quant_scheme import QuantizationScheme from torch.nn import Module, Parameter diff --git a/src/sparsetensors/quantization/observers/base.py b/src/sparsetensors/quantization/observers/base.py index 52a464b9..96fe1049 100644 --- a/src/sparsetensors/quantization/observers/base.py +++ b/src/sparsetensors/quantization/observers/base.py @@ -14,8 +14,8 @@ from typing import Optional, Tuple -from sparsetensors.quantization.quant_args import QuantizationArgs -from sparsetensors.registry.registry import RegistryMixin +from compressed_tensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.registry.registry import RegistryMixin from torch import FloatTensor, IntTensor, Tensor from torch.nn import Module diff --git a/src/sparsetensors/quantization/observers/helpers.py b/src/sparsetensors/quantization/observers/helpers.py index 0ec086a7..3fd2f4b6 100644 --- a/src/sparsetensors/quantization/observers/helpers.py +++ b/src/sparsetensors/quantization/observers/helpers.py @@ -15,7 +15,7 @@ from typing import Tuple import torch -from sparsetensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.quantization.quant_args import QuantizationArgs from torch import FloatTensor, IntTensor, Tensor diff --git a/src/sparsetensors/quantization/observers/memoryless.py b/src/sparsetensors/quantization/observers/memoryless.py index 0a70d6d2..04026807 100644 --- a/src/sparsetensors/quantization/observers/memoryless.py +++ b/src/sparsetensors/quantization/observers/memoryless.py @@ -15,8 +15,8 @@ from typing import Tuple import torch -from sparsetensors.quantization.observers.base import Observer -from sparsetensors.quantization.observers.helpers import calculate_qparams +from compressed_tensors.quantization.observers.base import Observer +from compressed_tensors.quantization.observers.helpers import calculate_qparams from torch import FloatTensor, IntTensor, Tensor diff --git a/src/sparsetensors/quantization/observers/min_max.py b/src/sparsetensors/quantization/observers/min_max.py index 1d2b4dc7..808f24c3 100644 --- a/src/sparsetensors/quantization/observers/min_max.py +++ b/src/sparsetensors/quantization/observers/min_max.py @@ -15,9 +15,9 @@ from typing import Tuple import torch -from sparsetensors.quantization.observers.base import Observer -from sparsetensors.quantization.observers.helpers import calculate_qparams -from sparsetensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.quantization.observers.base import Observer +from compressed_tensors.quantization.observers.helpers import calculate_qparams +from compressed_tensors.quantization.quant_args import QuantizationArgs from torch import FloatTensor, IntTensor, Tensor diff --git a/src/sparsetensors/quantization/quant_args.py b/src/sparsetensors/quantization/quant_args.py index 76bd61f0..64b5005f 100644 --- a/src/sparsetensors/quantization/quant_args.py +++ b/src/sparsetensors/quantization/quant_args.py @@ -80,6 +80,6 @@ def get_observer(self): """ :return: torch quantization FakeQuantize built based on these QuantizationArgs """ - from sparsetensors.quantization.observers.base import Observer + from compressed_tensors.quantization.observers.base import Observer return Observer.load_from_registry(self.observer, quantization_args=self) diff --git a/src/sparsetensors/quantization/quant_config.py b/src/sparsetensors/quantization/quant_config.py index 2a2b345f..7214bc83 100644 --- a/src/sparsetensors/quantization/quant_config.py +++ b/src/sparsetensors/quantization/quant_config.py @@ -16,8 +16,8 @@ from typing import Dict, List, Optional from pydantic import BaseModel, Field -from sparsetensors.quantization.quant_scheme import QuantizationScheme -from sparsetensors.quantization.utils import ( +from compressed_tensors.quantization.quant_scheme import QuantizationScheme +from compressed_tensors.quantization.utils import ( calculate_compression_ratio, is_module_quantized, iter_named_leaf_modules, diff --git a/src/sparsetensors/quantization/quant_scheme.py b/src/sparsetensors/quantization/quant_scheme.py index 7077c24e..c083fc55 100644 --- a/src/sparsetensors/quantization/quant_scheme.py +++ b/src/sparsetensors/quantization/quant_scheme.py @@ -15,7 +15,7 @@ from typing import List, Optional from pydantic import BaseModel -from sparsetensors.quantization.quant_args import QuantizationArgs +from compressed_tensors.quantization.quant_args import QuantizationArgs __all__ = ["QuantizationScheme"] diff --git a/src/sparsetensors/utils/helpers.py b/src/sparsetensors/utils/helpers.py index c584c2ee..ac9ed229 100644 --- a/src/sparsetensors/utils/helpers.py +++ b/src/sparsetensors/utils/helpers.py @@ -15,9 +15,9 @@ from typing import Optional -from sparsetensors.base import SPARSITY_CONFIG_NAME -from sparsetensors.compressors import ModelCompressor -from sparsetensors.config import CompressionConfig +from compressed_tensors.base import SPARSITY_CONFIG_NAME +from compressed_tensors.compressors import ModelCompressor +from compressed_tensors.config import CompressionConfig from transformers import AutoConfig diff --git a/tests/quantization/lifecycle/test_apply.py b/tests/quantization/lifecycle/test_apply.py index eeb29a41..6a3d17af 100644 --- a/tests/quantization/lifecycle/test_apply.py +++ b/tests/quantization/lifecycle/test_apply.py @@ -13,8 +13,8 @@ # limitations under the License. -from sparsetensors.quantization.lifecycle import apply_quantization_config -from sparsetensors.quantization.quant_config import ( +from compressed_tensors.quantization.lifecycle import apply_quantization_config +from compressed_tensors.quantization.quant_config import ( QuantizationConfig, QuantizationStatus, ) diff --git a/tests/quantization/test_quant_args.py b/tests/quantization/test_quant_args.py index c407eae5..46a4cc49 100644 --- a/tests/quantization/test_quant_args.py +++ b/tests/quantization/test_quant_args.py @@ -14,7 +14,7 @@ import pytest from pydantic import ValidationError -from sparsetensors.quantization import ( +from compressed_tensors.quantization import ( QuantizationArgs, QuantizationStrategy, QuantizationType, diff --git a/tests/quantization/test_quant_config.py b/tests/quantization/test_quant_config.py index 92b68ab7..68688c36 100644 --- a/tests/quantization/test_quant_config.py +++ b/tests/quantization/test_quant_config.py @@ -15,7 +15,7 @@ import pytest from pydantic import ValidationError -from sparsetensors.quantization import ( +from compressed_tensors.quantization import ( QuantizationConfig, QuantizationScheme, QuantizationStatus, diff --git a/tests/quantization/test_quant_scheme.py b/tests/quantization/test_quant_scheme.py index 63b135b5..1c198812 100644 --- a/tests/quantization/test_quant_scheme.py +++ b/tests/quantization/test_quant_scheme.py @@ -14,7 +14,7 @@ import pytest from pydantic import ValidationError -from sparsetensors.quantization import QuantizationArgs, QuantizationScheme +from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme def test_basic_scheme(): diff --git a/tests/test_bitmask.py b/tests/test_bitmask.py index b5bca142..28d29ed0 100644 --- a/tests/test_bitmask.py +++ b/tests/test_bitmask.py @@ -18,7 +18,7 @@ import pytest import torch from safetensors.torch import save_file -from sparsetensors import BitmaskCompressor, BitmaskConfig, BitmaskTensor +from compressed_tensors import BitmaskCompressor, BitmaskConfig, BitmaskTensor @pytest.mark.parametrize( diff --git a/tests/test_registry.py b/tests/test_registry.py index b73d357f..a183d77d 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -13,7 +13,7 @@ # limitations under the License. import pytest -from sparsetensors import ( +from compressed_tensors import ( BitmaskCompressor, BitmaskConfig, CompressionConfig,