Logger Module For Efficient Transformers

QEfficient/__init__.py

@@ -37,7 +37,9 @@
 from QEfficient.peft import QEffAutoPeftModelForCausalLM
 from QEfficient.transformers.transform import transform
 from QEfficient.utils import custom_format_warning
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA")
 
 # custom warning for the better logging experience
 warnings.formatwarning = custom_format_warning

QEfficient/base/modeling_qeff.py

@@ -7,7 +7,6 @@
 
 import gc
 import inspect
-import logging
 import shutil
 import subprocess
 import warnings
@@ -45,8 +44,9 @@
     to_named_specializations,
 )
 from QEfficient.utils.export_utils import export_wrapper
+from QEfficient.utils.logging_utils import QEFFLogger
 
-logger = logging.getLogger(__name__)
+logger = QEFFLogger.get_logger("INFRA")
 
 
 class QEFFBaseModel(ABC):
@@ -382,6 +382,7 @@ def _export(
             raise e
 
         self.onnx_path = onnx_path
+        logger.info("Model export is finished and saved at: %s", onnx_path)
         return onnx_path
 
     def get_onnx_path(
@@ -686,4 +687,5 @@ def _compile(
         logger.info("Hashed parameters exported successfully.")
 
         self.qpc_path = qpc_path
+        logger.info("Model compilation is finished and saved at: %s", qpc_path)
         return qpc_path

QEfficient/base/pytorch_transforms.py

@@ -9,7 +9,9 @@
 
 from torch import nn
 
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA")
 
 
 class PytorchTransform:

QEfficient/cloud/export.py

@@ -12,7 +12,9 @@
 from QEfficient.base.common import QEFFCommonLoader
 from QEfficient.utils import check_and_assign_cache_dir
 from QEfficient.utils.custom_yaml import generate_custom_io
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA")
 
 # 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("INFRA")
 
 
 # TODO: Remove after adding support for VLM's compile and execute

QEfficient/compile/compile_helper.py

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

QEfficient/compile/qnn_compiler.py

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

QEfficient/diffusers/models/transformers/transformer_flux.py

@@ -20,7 +20,9 @@
 )
 
 from QEfficient.diffusers.models.modeling_utils import compute_blocked_attention, get_attention_blocking_config
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 def qeff_apply_rotary_emb(

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -38,7 +38,9 @@
     set_execute_params,
 )
 from QEfficient.generation.cloud_infer import QAICInferenceSession
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 class QEffFluxPipeline:

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -16,7 +16,9 @@
 from tqdm import tqdm
 
 from QEfficient.utils._utils import load_json
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 def calculate_compressed_latent_dimension(height: int, width: int, vae_scale_factor: int) -> int:

QEfficient/diffusers/pipelines/wan/pipeline_wan.py

@@ -37,7 +37,9 @@
 )
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 class QEffWanPipeline:

QEfficient/diffusers/pipelines/wan/pipeline_wan_i2v.py

@@ -41,7 +41,9 @@
 )
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 class QEffWanImageToVideoPipeline:

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("INFRA")
 
 
 def convert_to_cloud_bertstyle(

QEfficient/generation/embedding_handler.py

@@ -23,7 +23,9 @@
 
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA")
 
 
 class VisionHandler:

QEfficient/generation/text_generation_inference.py

@@ -19,9 +19,11 @@
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 from QEfficient.utils import padding_check_and_fix
 from QEfficient.utils.constants import Constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
 from QEfficient.utils.sampler_utils import validate_sampler_inputs
 
+logger = QEFFLogger.get_logger("INFRA")
+
 
 @dataclass
 class PerfMetrics:
@@ -1326,4 +1328,6 @@ def generate(
             generated_ids=self._qaic_model.generated_ids,
             perf_metrics=perf_metrics,
         )
+        logger.info("Text generation finished")
+        QEFFLogger.print_table()
         return latency_stats

QEfficient/generation/vlm_generation.py

@@ -37,7 +37,9 @@
 )
 from QEfficient.utils import LRUCache
 from QEfficient.utils.constants import Constants
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA")
 
 
 class VisionLanguageGeneration(QEffTextGenerationBase):

QEfficient/peft/auto.py

@@ -6,7 +6,6 @@
 # ----------------------------------------------------------------------------
 
 import hashlib
-import logging
 import warnings
 from typing import List, Optional, Union
 
@@ -32,8 +31,9 @@
 from QEfficient.utils import constants
 from QEfficient.utils._utils import get_padding_shape_from_config
 from QEfficient.utils.hash_utils import to_hashable
+from QEfficient.utils.logging_utils import QEFFLogger
 
-logger = logging.getLogger(__name__)
+logger = QEFFLogger.get_logger("FT")
 
 
 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.hash_utils import to_hashable
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT")
 
 
 class QEffAutoLoraModelForCausalLM(QEFFAutoModelForCausalLM):

QEfficient/transformers/modeling_utils.py

@@ -92,7 +92,7 @@
 from QEfficient.customop import CustomRMSNormAIC
 from QEfficient.proxy.pytorch_transform import QeffProxyModuleTransform
 from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
 
 if TYPE_CHECKING:
     from QEfficient.base.modeling_qeff import QEFFBaseModel
@@ -164,6 +164,8 @@
     QEffWhisperPositionalEmbedding,
 )
 
+logger = QEFFLogger.get_logger("MODEL")
+
 # Define a named tuple for ModelArchitectures
 # Required for the Automation tool
 ModelArchitectures = namedtuple("ModelArchitectures", ["architectures"])

QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -39,7 +39,9 @@
 from QEfficient.transformers.cache_utils import QEffHybridCacheForGPTOSS
 from QEfficient.transformers.modeling_attn_mask_utils import _create_causal_mask
 from QEfficient.utils.constants import MIN_MASKED_ATTENTION_VALUE
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("MODEL")
 
 
 class QEffGptOssExperts(GptOssExperts):

QEfficient/transformers/models/internvl/modeling_internvl.py

@@ -13,7 +13,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("MODEL")
 
 
 class QEffInternEncoderWrapper(nn.Module):

QEfficient/transformers/models/llava/modeling_llava.py

@@ -15,7 +15,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("MODEL")
 
 BS = 1
 FBS = 4

QEfficient/transformers/models/llava_next/modeling_llava_next.py

@@ -18,7 +18,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("MODEL")
 
 BS = constants.ONNX_EXPORT_EXAMPLE_BATCH_SIZE
 FBS = constants.ONNX_EXPORT_EXAMPLE_FBS

QEfficient/transformers/models/mistral3/modeling_mistral3.py

@@ -21,7 +21,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("MODEL")
 
 
 def custom_cumsum(tensor):

QEfficient/transformers/models/modeling_auto.py

@@ -76,9 +76,10 @@
     get_padding_shape_from_config,
 )
 from QEfficient.utils.check_ccl_specializations import process_ccl_specializations
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
 from QEfficient.utils.sampler_utils import get_sampling_inputs_and_outputs
 
+logger = QEFFLogger.get_logger("MODEL")
 CUSTOM_IO_DTYPE_MAP = {
     torch.float16: "float16",
     torch.bfloat16: "bfloat16",
@@ -153,7 +154,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, *args, **kwargs):
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
-
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
 
         kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {})
@@ -315,7 +316,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
-
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
 
         # This is support models that should be classified to in a different auto class but transformers load them via this class
@@ -695,7 +696,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
-
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
         kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {})
         return cls(model, pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
@@ -1271,7 +1272,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, qaic_config: Option
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
-
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
 
         kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {})
@@ -2085,6 +2086,7 @@ def from_pretrained(
         config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
         config._attn_implementation = "eager"
         config.vision_config.use_flash_attn = "false"
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, config, *args, **kwargs)
 
         kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {})
@@ -2697,6 +2699,7 @@ def from_pretrained(
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
 
         kwargs.update({"enable_proxy": enable_proxy} if enable_proxy else {})
@@ -2949,6 +2952,7 @@ def from_pretrained(
         kv_offload = kwargs.pop("kv_offload", None)
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
         if qaic_config is not None:
             qaic_config["pretrained_model_name_or_path"] = pretrained_model_name_or_path
@@ -4209,7 +4213,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k
             logger.warning("Updating low_cpu_mem_usage=False")
 
         kwargs.update({"attn_implementation": "eager", "low_cpu_mem_usage": False})
-
+        logger.info("Initiating the model weight loading.")
         model = cls._hf_auto_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
 
         # This is support models that should be classified to in a different auto class but transformers load them via this class