Skip to content

Commit

Permalink
Add: targets and ignore to sparsity compression config (#159)
Browse files Browse the repository at this point in the history
  • Loading branch information
rahul-tuli authored Sep 20, 2024
1 parent a654f6c commit 83b2e7a
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion src/compressed_tensors/config/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

from enum import Enum
from typing import Optional
from typing import List, Optional

from compressed_tensors.registry import RegistryMixin
from pydantic import BaseModel
Expand All @@ -37,11 +37,16 @@ class SparsityCompressionConfig(RegistryMixin, BaseModel):
Base data class for storing sparsity compression parameters
:param format: name of compression format
:param targets: List of layer names or layer types that aren't sparse and should
be ignored during compression. By default, assume all layers are targeted
:param ignore: List of layer names (unique) to ignore from targets. Defaults to None
:param global_sparsity: average sparsity of the entire model
:param sparsity_structure: structure of the sparsity, such as
"unstructured", "2:4", "8:16" etc
"""

format: str
targets: Optional[List[str]] = None
ignore: Optional[List[str]] = None
global_sparsity: Optional[float] = 0.0
sparsity_structure: Optional[str] = "unstructured"

0 comments on commit 83b2e7a

Please sign in to comment.