Logger module in Efficient Transformers

QEfficient/__init__.py

@@ -5,6 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
+import atexit
 import os
 import warnings
 
@@ -16,7 +17,12 @@
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 # Placeholder for all non-transformer models registered in QEfficient
 import QEfficient.utils.model_registery  # noqa: F401
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
+
+# Register the cleanup function
+atexit.register(QEFFLogger.close_logger)
 
 # custom warning for the better logging experience
 warnings.formatwarning = custom_format_warning

QEfficient/base/modeling_qeff.py

@@ -7,7 +7,6 @@
 
 import hashlib
 import inspect
-import logging
 import shutil
 import subprocess
 import warnings
@@ -24,8 +23,9 @@
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants, create_json, dump_qconfig, generate_mdp_partition_config, load_json
 from QEfficient.utils.cache import QEFF_HOME, to_hashable
+from QEfficient.utils.logging_utils import QEFFLogger
 
-logger = logging.getLogger(__name__)
+logger = QEFFLogger.get_logger()
 
 
 class QEFFBaseModel(ABC):
@@ -63,6 +63,11 @@ def __init__(self, model: torch.nn.Module) -> None:
         else:
             logger.info(f"Pytorch transforms applied to model: {self.model_name}")
 
+    def set_loglevel(loglevel: Optional[str] = "INFO"):
+        QEFFLogger.set_loglevel(loglevel)
+        logger.info(f"log level to be set is: {loglevel}")
+        logger.debug(f"log level set to: {loglevel}")
+
     @property
     @abstractmethod
     def model_name(self) -> str: ...

QEfficient/cloud/export.py

@@ -11,7 +11,9 @@
 
 from QEfficient.base.common import QEFFCommonLoader
 from QEfficient.utils import check_and_assign_cache_dir
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 # Specifically for Docker images.
 ROOT_DIR = os.path.dirname(os.path.abspath(""))

QEfficient/cloud/infer.py

@@ -17,7 +17,9 @@
 
 from QEfficient.base.common import QEFFCommonLoader
 from QEfficient.utils import check_and_assign_cache_dir, load_hf_processor, load_hf_tokenizer
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 # TODO: Remove after adding support for VLM's compile and execute

QEfficient/compile/compile_helper.py

@@ -14,7 +14,9 @@
 
 from QEfficient.compile.qnn_compiler import compile as qnn_compile
 from QEfficient.utils._utils import load_json, load_yaml
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 def create_and_dump_specializations(

QEfficient/compile/qnn_compiler.py

@@ -18,7 +18,9 @@
     generate_data_format_config,
     generate_qnn_specialization,
 )
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QNN:

QEfficient/exporter/export_hf_to_cloud_ai_100.py

@@ -20,7 +20,9 @@
 from QEfficient.utils import load_hf_tokenizer
 from QEfficient.utils.constants import QEFF_MODELS_DIR, Constants
 from QEfficient.utils.generate_inputs import InputHandler
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 def convert_to_cloud_bertstyle(

QEfficient/generation/text_generation_inference.py

@@ -18,7 +18,9 @@
 
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import padding_check_and_fix
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 @dataclass

QEfficient/peft/auto.py

@@ -6,7 +6,6 @@
 # ----------------------------------------------------------------------------
 
 import hashlib
-import logging
 import warnings
 from typing import List, Optional, Union
 
@@ -28,8 +27,9 @@
 from QEfficient.utils import constants
 from QEfficient.utils._utils import get_padding_shape_from_config
 from QEfficient.utils.cache import to_hashable
+from QEfficient.utils.logging_utils import QEFFLogger
 
-logger = logging.getLogger(__name__)
+logger = QEFFLogger.get_logger()
 
 
 class QEffAutoPeftModelForCausalLM(QEFFBaseModel):

QEfficient/peft/lora/auto.py

@@ -19,7 +19,9 @@
 from QEfficient.peft.lora.pytorch_transforms import LoraModelInputsTransform, TargetModulesTransform
 from QEfficient.utils import constants, get_padding_shape_from_config
 from QEfficient.utils.cache import to_hashable
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffAutoLoraModelForCausalLM(QEFFAutoModelForCausalLM):

QEfficient/transformers/models/granitemoe/modeling_granitemoe.py

@@ -22,13 +22,15 @@
     GraniteMoeRotaryEmbedding,
     GraniteMoeTopKGating,
     load_balancing_loss_func,
-    logger,
     repeat_kv,
     rotate_half,
 )
 
 from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffGraniteMoeRotaryEmbedding(GraniteMoeRotaryEmbedding):
@@ -199,9 +201,7 @@ def forward(
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if self.gradient_checkpointing and self.training and use_cache:
-            logger.warning_once(
-                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
-            )
+            logger.warning("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.")
             use_cache = False
 
         if inputs_embeds is None:
@@ -216,7 +216,7 @@ def forward(
                 past_key_values = QEffDynamicCache()
             else:
                 past_key_values = QEffDynamicCache.from_legacy_cache(past_key_values)
-                logger.warning_once(
+                logger.warning(
                     "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
                     "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
                     "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"

QEfficient/transformers/models/internvl/modeling_internvl.py

@@ -11,7 +11,9 @@
 
 from QEfficient.utils import constants
 from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffInternEncoderWrapper(nn.Module):

QEfficient/transformers/models/llama_swiftkv/modeling_llama_swiftkv.py

@@ -20,7 +20,7 @@
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
-from transformers.models.llama.modeling_llama import LlamaMLP, LlamaRMSNorm, logger, repeat_kv
+from transformers.models.llama.modeling_llama import LlamaMLP, LlamaRMSNorm, repeat_kv
 
 from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
@@ -29,6 +29,9 @@
     QEffLlamaRotaryEmbedding,
     qeff_apply_rotary_pos_emb,
 )
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffLlamaSwiftKVConfig(LlamaConfig):
@@ -302,7 +305,7 @@ def forward(
                 past_key_values = QEffDynamicCache()
             else:
                 past_key_values = QEffDynamicCache.from_legacy_cache(past_key_values)
-                logger.warning_once(
+                logger.warning(
                     "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
                     "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
                     "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"

QEfficient/transformers/models/llava/modeling_llava.py

@@ -13,7 +13,9 @@
 )
 
 from QEfficient.utils._utils import IOInfo
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 BS = 1
 NUM_CHANNEL = 3

QEfficient/transformers/models/llava_next/modeling_llava_next.py

@@ -16,7 +16,9 @@
 
 from QEfficient.utils import constants
 from QEfficient.utils._utils import IOInfo
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffLlavaNextEncoderWrapper(nn.Module):
@@ -65,7 +67,7 @@ def forward(self, pixel_values, image_sizes):
                     np.prod(image_feature.shape) % (num_patch_height * num_patch_width * height * width) != 0
                     and vision_feature_select_strategy == "default"
                 ):
-                    logger.warning_once(
+                    logger.warning(
                         "Image feature shape does not line up with the provided patch size. "
                         "You may be using the `default` vision_feature_select_strategy with a"
                         " visual encoder that does not have CLS."

QEfficient/transformers/models/mistral/modeling_mistral.py

@@ -24,13 +24,15 @@
     MistralForCausalLM,
     MistralModel,
     MistralRotaryEmbedding,
-    logger,
     repeat_kv,
     rotate_half,
 )
 
 from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffMistralRotaryEmbedding(MistralRotaryEmbedding):

QEfficient/transformers/models/mllama/modeling_mllama.py

@@ -30,7 +30,6 @@
     MllamaTextModel,
     MllamaTextSelfAttention,
     MllamaVisionModel,
-    logger,
     repeat_kv,
     rotate_half,
 )
@@ -43,6 +42,9 @@
 )
 from QEfficient.utils import constants
 from QEfficient.utils._utils import IOInfo
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 MAX_NUM_IMG = 1
 NUM_CHANNEL = 3

QEfficient/transformers/models/modeling_auto.py

@@ -57,7 +57,9 @@
     get_padding_shape_from_config,
 )
 from QEfficient.utils.cache import to_hashable
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEFFTransformersBase(QEFFBaseModel):

QEfficient/transformers/models/whisper/modeling_whisper.py

@@ -24,12 +24,14 @@
     WhisperForConditionalGeneration,
     WhisperModel,
     WhisperPositionalEmbedding,
-    logger,
 )
 
 from QEfficient.transformers.cache_utils import QEffEncoderDecoderCache
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
 from QEfficient.utils._utils import IOInfo
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffWhisperPositionalEmbedding(WhisperPositionalEmbedding):

QEfficient/transformers/quantizers/quantizer_awq.py

@@ -15,7 +15,9 @@
     replace_linear_layer_with_target_layer,
     replace_quantization_scales,
 )
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffAwqConfig(AwqConfig):

QEfficient/transformers/quantizers/quantizer_compressed_tensors.py

@@ -14,7 +14,9 @@
 from transformers.utils.quantization_config import CompressedTensorsConfig, QuantizationConfigMixin, QuantizationMethod
 
 from QEfficient.transformers.quantizers.quantizer_utils import get_keys_to_not_convert
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 FP8_DTYPE = torch.float8_e4m3fn
 

QEfficient/transformers/quantizers/quantizer_gptq.py

@@ -15,7 +15,9 @@
     repack_zeros,
     replace_linear_layer_with_target_layer,
 )
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class QEffGPTQConfig(GPTQConfig):

QEfficient/transformers/transform.py

@@ -13,7 +13,9 @@
 from QEfficient.base.modeling_qeff import QEFFBaseModel
 from QEfficient.transformers.cache_utils import QEffDynamicCache
 from QEfficient.transformers.modeling_utils import TransformersToQEffModulesDict
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 def replace_module_with_qeff_layers(model: nn.Module) -> None:

QEfficient/utils/_utils.py

@@ -26,7 +26,9 @@
 )
 
 from QEfficient.utils.constants import QEFF_MODELS_DIR, Constants, QnnConstants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 class DownloadRetryLimitExceeded(Exception):

QEfficient/utils/device_utils.py

@@ -10,7 +10,9 @@
 import subprocess
 
 from QEfficient.utils.constants import Constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger()
 
 
 def is_networks_loaded(stdout):