basic support

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/__init__.py

@@ -17,6 +17,7 @@
 # flake8: noqa
 from .compressors import *
 from .config import *
+from .logger import LoggerConfig, configure_logger, logger
 from .quantization import QuantizationConfig, QuantizationStatus
 from .utils import *
 from .version import *

src/compressed_tensors/logger.py

@@ -0,0 +1,132 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Logger configuration for Compressed Tensors.
+"""
+
+import os
+import sys
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from loguru import logger
+
+
+__all__ = ["LoggerConfig", "configure_logger", "logger"]
+
+
+# used by `support_log_once``
+_logged_once = set()
+
+
+@dataclass
+class LoggerConfig:
+    disabled: bool = False
+    clear_loggers: bool = True
+    console_log_level: Optional[str] = "INFO"
+    log_file: Optional[str] = None
+    log_file_level: Optional[str] = None
+
+
+def configure_logger(config: Optional[LoggerConfig] = None) -> None:
+    """
+    Configure the logger for Compressed Tensors.
+    This function sets up the console and file logging
+    as per the specified or default parameters.
+
+    Note: Environment variables take precedence over the function parameters.
+
+    :param config: The configuration for the logger to use.
+    :type config: LoggerConfig
+    """
+    logger_config = config or LoggerConfig()
+
+    # env vars get priority
+    if (disabled := os.getenv("COMPRESSED_TENSORS_LOG_DISABLED")) is not None:
+        logger_config.disabled = disabled.lower() == "true"
+    if (clear_loggers := os.getenv("COMPRESSED_TENSORS_CLEAR_LOGGERS")) is not None:
+        logger_config.clear_loggers = clear_loggers.lower() == "true"
+    if (console_log_level := os.getenv("COMPRESSED_TENSORS_LOG_LEVEL")) is not None:
+        logger_config.console_log_level = console_log_level.upper()
+    if (log_file := os.getenv("COMPRESSED_TENSORS_LOG_FILE")) is not None:
+        logger_config.log_file = log_file
+    if (log_file_level := os.getenv("COMPRESSED_TENSORS_LOG_FILE_LEVEL")) is not None:
+        logger_config.log_file_level = log_file_level.upper()
+
+    if logger_config.disabled:
+        logger.disable("compressed_tensors")
+        return
+
+    logger.enable("compressed_tensors")
+
+    if logger_config.clear_loggers:
+        logger.remove()
+
+    if logger_config.console_log_level:
+        # log as a human readable string with the time, function, level, and message
+        logger.add(
+            sys.stdout,
+            level=logger_config.console_log_level.upper(),
+            format="{time} | {function} | {level} - {message}",
+            filter=support_log_once,
+        )
+
+    if logger_config.log_file or logger_config.log_file_level:
+        log_file = logger_config.log_file or "compressed_tensors.log"
+        log_file_level = logger_config.log_file_level or "INFO"
+        # log as json to the file for easier parsing
+        logger.add(
+            log_file,
+            level=log_file_level.upper(),
+            serialize=True,
+            filter=support_log_once,
+        )
+
+
+def support_log_once(record: Dict[str, Any]) -> bool:
+    """
+    Support logging only once using `.bind(log_once=True)`
+
+    ```
+    logger.bind(log_once=False).info("This will log multiple times")
+    logger.bind(log_once=False).info("This will log multiple times")
+    logger.bind(log_once=True).info("This will only log once")
+    logger.bind(log_once=True).info("This will only log once")  # skipped
+    ```
+    """
+    log_once = record["extra"].get("log_once", False)
+    level = getattr(record["level"], "name", "none")
+    message = str(level) + record["message"]
+
+    if log_once and message in _logged_once:
+        return False
+
+    if log_once:
+        _logged_once.add(message)
+
+    return True
+
+
+# invoke logger setup on import with default values enabling console logging with INFO
+# and disabling file logging
+configure_logger(
+    config=LoggerConfig(
+        disabled=False,
+        clear_loggers=True,
+        console_log_level="INFO",
+        log_file=None,
+        log_file_level=None,
+    )
+)

src/compressed_tensors/quantization/lifecycle/forward.py

@@ -29,7 +29,6 @@
     calculate_range,
     compute_dynamic_scales_and_zp,
 )
-from compressed_tensors.utils import safe_permute
 from torch.nn import Module
 
 
@@ -265,8 +264,7 @@ def _process_quantization(
     ):
 
         output_dtype = dtype if dtype is not None else x.dtype
-        output = torch.zeros_like(x).to(output_dtype)
-        columns = output.shape[-1]
+        columns = x.size(-1)
 
         # TODO: make validation step for inputs
 
@@ -294,7 +292,7 @@ def _process_quantization(
             group_sizes = group_sizes[torch.argsort(group_indices)]
 
             perm = torch.argsort(g_idx)
-            x = safe_permute(x, perm, dim=1)
+            x = x.index_select(dim=-1, index=perm)
 
         # Maintain all dimensions except the last dim, which is divided by group_size
         reshaped_dims = (
@@ -324,11 +322,11 @@ def _process_quantization(
                 global_scale=global_scale,
             )
 
-        output = output.flatten(start_dim=-2)
+        output = output.flatten(-2, -1)
         output = output.to(output_dtype)
 
         if not is_column_order:
-            output = safe_permute(output, torch.argsort(perm), dim=1)
+            output = output.index_select(dim=-1, index=torch.argsort(perm))
 
     else:  # covers channel, token and tensor strategies
         if do_quantize:

src/compressed_tensors/utils/helpers.py

@@ -15,14 +15,18 @@
 import contextlib
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, TypeVar
 
 import numpy
 import torch
 from frozendict import frozendict
+from loguru import logger
 from transformers import AutoConfig
 
 
+T = TypeVar("T", bound="Callable")  # used by `deprecated`
+
+
 if TYPE_CHECKING:
     from compressed_tensors.compressors import ModelCompressor
 
@@ -170,15 +174,17 @@ def getattr_chain(obj: Any, chain_str: str, *args, **kwargs) -> Any:
     return res
 
 
-def deprecated(future_name: Optional[str] = None, message: Optional[str] = None):
+def deprecated(
+    future_name: Optional[str] = None, message: Optional[str] = None
+) -> Callable[[T], T]:
     """
     Decorator to mark functions as deprecated
 
     :param new_function: Function called in place of deprecated function
     :param message: Deprecation message, replaces default deprecation message
     """
 
-    def decorator(func: Callable[[Any], Any]):
+    def decorator(func: T) -> T:
         nonlocal message
 
         if message is None:
@@ -190,7 +196,7 @@ def decorator(func: Callable[[Any], Any]):
 
         @wraps(func)
         def wrapped(*args, **kwargs):
-            warnings.warn(message, DeprecationWarning, stacklevel=2)
+            logger.bind(log_once=True).warning(message)
             return func(*args, **kwargs)
 
         return wrapped

src/compressed_tensors/utils/permute.py

@@ -12,18 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Set, Tuple
-
 import torch
+from compressed_tensors.utils.helpers import deprecated
 
 
 __all__ = ["safe_permute"]
 
 
-# these datatypes are missing implementations required for standard permutation
-_EXPERIMENTAL_DTYPES: Set[Tuple[torch.dtype, torch.device]] = set()
-
-
+@deprecated("Tensor.index_select")
 def safe_permute(value: torch.Tensor, perm: torch.Tensor, dim: int = 0) -> torch.Tensor:
     """
     Perform out-of-place permutation without using torch.Tensor.index_put_,
@@ -34,37 +30,4 @@ def safe_permute(value: torch.Tensor, perm: torch.Tensor, dim: int = 0) -> torch
     :param dim: dimension along which to apply permutation
     :return: permuted value
     """
-    dtype_tuple = (value.dtype, value.device)
-
-    if dtype_tuple in _EXPERIMENTAL_DTYPES:
-        return _fallback_permute(value, perm, dim)
-
-    try:
-        return value[tuple([slice(None)] * dim + [perm])]
-    except RuntimeError:
-        # Mark dtype as experimental if advanced indexing fails
-        _EXPERIMENTAL_DTYPES.add(dtype_tuple)
-        return _fallback_permute(value, perm, dim)
-
-
-def _fallback_permute(
-    value: torch.Tensor, perm: torch.Tensor, dim: int
-) -> torch.Tensor:
-    """
-    Fallback permutation method for experimental dtypes.
-
-    :param value: tensor to permute
-    :param perm: permutation map
-    :param dim: dimension along which to apply permutation
-    :return: permuted value
-    """
-    value_ret = value.clone()  # cannot use zeros_like b/c of missing impl.
-    orig_slices = [slice(None)] * (dim + 1)
-    perm_slices = [slice(None)] * (dim + 1)
-
-    for index, perm_index in enumerate(perm):
-        orig_slices[dim] = index
-        perm_slices[dim] = perm_index
-        value_ret[tuple(orig_slices)] = value[tuple(perm_slices)]
-
-    return value_ret
+    return value.index_select(dim, perm)