Logger Module For Efficient Transformers

QEfficient/__init__.py

@@ -35,7 +35,17 @@
 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", loglevel="INFO")
+
+# custom warning for the better logging experience
+warnings.formatwarning = custom_format_warning
+
+
+# custom warning for the better logging experience
+warnings.formatwarning = custom_format_warning
+
 
 # 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
@@ -35,8 +34,9 @@
     load_json,
 )
 from QEfficient.utils.export_utils import export_wrapper
+from QEfficient.utils.logging_utils import QEFFLogger
 
-logger = logging.getLogger(__name__)
+logger = QEFFLogger.get_logger("INFRA", loglevel="INFO")
 
 
 class QEFFBaseModel(ABC):
@@ -326,6 +326,7 @@ def _export(
             self.prefill_onnx_path = onnx_path
         else:
             self.onnx_path = onnx_path
+        logger.info("Model export is finished and saved at: %s", onnx_path)
         return onnx_path
 
     def get_onnx_path(
@@ -539,4 +540,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", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 # Specifically for Docker images.
 ROOT_DIR = os.path.dirname(os.path.abspath(""))

QEfficient/cloud/finetune.py

@@ -29,10 +29,12 @@
 from QEfficient.finetune.utils.dataset_utils import get_dataloader, get_longest_seq_length
 from QEfficient.finetune.utils.device_map import get_device_map
 from QEfficient.finetune.utils.helper import Task_Mode, get_world_size
-from QEfficient.finetune.utils.logging_utils import logger
 from QEfficient.finetune.utils.parser import get_finetune_parser
 from QEfficient.finetune.utils.train_utils import print_model_size, print_trainable_parameters, train
 from QEfficient.utils._utils import hf_download
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 # Try importing QAIC-specific module, proceed without it if unavailable
 try:

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", loglevel="INFO")
 
 
 # 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
-from QEfficient.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("INFRA", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 
 class QNN:

QEfficient/diffusers/models/transformers/transformer_flux.py

@@ -18,8 +18,9 @@
     _get_qkv_projections,
 )
 
-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", loglevel="INFO")
 
 
 def qeff_apply_rotary_emb(
@@ -90,9 +91,9 @@
             key = qeff_apply_rotary_emb(key, image_rotary_emb)
 
         # Get blocking configuration
         blocking_mode, head_block_size, num_kv_blocks, num_q_blocks = get_attention_blocking_config()
         # Apply blocking using pipeline_utils
         hidden_states = compute_blocked_attention(
             query.transpose(1, 2),
             key.transpose(1, 2),
             value.transpose(1, 2),

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -38,7 +38,9 @@
     set_module_device_ids,
 )
 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", loglevel="INFO")
 
 
 class QEffFluxPipeline:

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -18,7 +18,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", loglevel="INFO")
 
 
 def calculate_compressed_latent_dimension(height: int, width: int, vae_scale_factor: int) -> int:

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", loglevel="INFO")
 
 
 def convert_to_cloud_bertstyle(

QEfficient/finetune/dataset/alpaca_dataset.py

@@ -11,7 +11,9 @@
 import torch
 from torch.utils.data import Dataset
 
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 PROMPT_DICT = {
     "prompt_input": (

QEfficient/finetune/dataset/custom_dataset.py

@@ -9,7 +9,9 @@
 import logging
 from pathlib import Path
 
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 
 def load_module_from_py_file(py_file: str) -> object:

QEfficient/finetune/dataset/grammar_dataset.py

@@ -10,7 +10,9 @@
 from datasets import load_dataset
 from torch.utils.data import Dataset
 
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 
 class grammar(Dataset):

QEfficient/finetune/utils/config_utils.py

@@ -20,7 +20,9 @@
 from QEfficient.finetune.configs.training import TrainConfig
 from QEfficient.finetune.dataset.dataset_config import DATASET_PREPROC
 from QEfficient.finetune.utils.helper import Peft_Method
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 
 def update_config(config, **kwargs):

QEfficient/finetune/utils/dataset_utils.py

@@ -16,7 +16,9 @@
 from QEfficient.finetune.data.sampler import DistributedLengthBasedBatchSampler
 from QEfficient.finetune.dataset.dataset_config import DATALOADER_COLLATE_FUNC, DATASET_PREPROC
 from QEfficient.finetune.utils.helper import get_world_size
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 
 def get_preprocessed_dataset(

QEfficient/finetune/utils/plot_metrics.py

@@ -11,7 +11,9 @@
 
 import matplotlib.pyplot as plt
 
-from QEfficient.finetune.utils.logging_utils import logger
+from QEfficient.utils.logging_utils import QEFFLogger
+
+logger = QEFFLogger.get_logger("FT", loglevel="INFO")
 
 
 def plot_metric(data, metric_name, x_label, y_label, title, colors):

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", loglevel="INFO")
 
 
 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", loglevel="INFO")
+
 
 @dataclass
 class PerfMetrics:
@@ -1316,4 +1318,5 @@ def generate(
             generated_ids=self._qaic_model.generated_ids,
             perf_metrics=perf_metrics,
         )
+        logger.info("Text Generated finised")
         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", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 
 class QEffAutoLoraModelForCausalLM(QEFFAutoModelForCausalLM):

QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -33,7 +33,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", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 
 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", loglevel="INFO")
 
 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", loglevel="INFO")
 
 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", loglevel="INFO")
 
 
 def custom_cumsum(tensor):

QEfficient/transformers/models/modeling_auto.py

@@ -69,9 +69,11 @@
     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", loglevel="INFO")
+
 
 class QEFFTransformersBase(QEFFBaseModel):
     """

QEfficient/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py

@@ -38,7 +38,9 @@
 from QEfficient.utils import constants
 from QEfficient.utils._utils import IOInfo, get_padding_shape_from_config
 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", loglevel="INFO")
 
 
 def qeff_apply_rotary_pos_emb(q, k, cos, sin, position_ids, mrope_section, unsqueeze_dim=1):

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("MODEL", loglevel="INFO")
 
 
 class QEffAwqConfig(AwqConfig):