Update logging_utils and log for zero rank

Signed-off-by: Mamta Singh <[email protected]>

Author: quic-mamta

QEfficient/cloud/finetune.py

@@ -34,13 +34,13 @@
 )
 from QEfficient.finetune.utils.train_utils import get_longest_seq_length, print_model_size, train
 from QEfficient.utils._utils import login_and_download_hf_lm
-from QEfficient.utils.logging_utils import ft_logger as logger
+from QEfficient.utils.logging_utils import logger
 
 # Try importing QAIC-specific module, proceed without it if unavailable
 try:
     import torch_qaic  # noqa: F401
 except ImportError as e:
-    logger.warning(f"{e}. Moving ahead without these qaic modules.")
+    logger.log_rank_zero(f"{e}. Moving ahead without these qaic modules.")
 
 logger.setLevel(logging.INFO)
 
@@ -121,7 +121,7 @@ def load_model_and_tokenizer(
         )
 
         if not hasattr(model, "base_model_prefix"):
-            raise RuntimeError("Given huggingface model does not have 'base_model_prefix' attribute.")
+            logger.raise_runtimeerror("Given huggingface model does not have 'base_model_prefix' attribute.")
 
         for param in getattr(model, model.base_model_prefix).parameters():
             param.requires_grad = False
@@ -146,7 +146,7 @@ def load_model_and_tokenizer(
     # If there is a mismatch between tokenizer vocab size and embedding matrix,
     # throw a warning and then expand the embedding matrix
     if len(tokenizer) > model.get_input_embeddings().weight.shape[0]:
-        logger.warning("Resizing the embedding matrix to match the tokenizer vocab size.")
+        logger.log_rank_zero("Resizing the embedding matrix to match the tokenizer vocab size.", logger.WARNING)
         model.resize_token_embeddings(len(tokenizer))
 
     # FIXME (Meet): Cover below line inside the logger once it is implemented.
@@ -162,7 +162,9 @@ def load_model_and_tokenizer(
         if hasattr(model, "supports_gradient_checkpointing") and model.supports_gradient_checkpointing:
             model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"preserve_rng_state": False})
         else:
-            raise RuntimeError("Given model doesn't support gradient checkpointing. Please disable it and run it.")
+            logger.raise_runtimeerror(
+                "Given model doesn't support gradient checkpointing. Please disable it and run it."
+            )
 
     model = apply_peft(model, train_config, peft_config_file, **kwargs)
 
@@ -222,7 +224,7 @@ def setup_dataloaders(
             - Length of longest sequence in the dataset.
 
     Raises:
-        ValueError: If validation is enabled but the validation set is too small.
+        RuntimeError: If validation is enabled but the validation set is too small.
 
     Notes:
         - Applies a custom data collator if provided by get_custom_data_collator.
@@ -246,12 +248,12 @@ def setup_dataloaders(
     #         )
     ##
     train_dl_kwargs = get_dataloader_kwargs(train_config, dataset_train, dataset_processer, "train")
-    logger.info(f"length of dataset_train = {len(dataset_train)}")
+    logger.log_rank_zero(f"Length of dataset_train = {len(dataset_train)}")
 
     # FIXME (Meet): Add custom data collator registration from the outside by the user.
     custom_data_collator = get_custom_data_collator(dataset_processer, dataset_config)
     if custom_data_collator:
-        logger.info("custom_data_collator is used")
+        logger.log_rank_zero("Custom_data_collator is used")
         train_dl_kwargs["collate_fn"] = custom_data_collator
 
     # Create DataLoaders for the training and validation dataset
@@ -261,7 +263,7 @@ def setup_dataloaders(
         pin_memory=True,
         **train_dl_kwargs,
     )
-    logger.info(f"Num of Training Set Batches loaded = {len(train_dataloader)}")
+    logger.log_rank_zero(f"Number of Training Set Batches loaded = {len(train_dataloader)}")
 
     eval_dataloader = None
     if train_config.run_validation:
@@ -281,11 +283,11 @@ def setup_dataloaders(
             **val_dl_kwargs,
         )
         if len(eval_dataloader) == 0:
-            raise ValueError(
+            logger.raise_runtimeerror(
                 f"The eval set size is too small for dataloader to load even one batch. Please increase the size of eval set. ({len(eval_dataloader)=})"
             )
         else:
-            logger.info(f"Num of Validation Set Batches loaded = {len(eval_dataloader)}")
+            logger.log_rank_zero(f"Number of Validation Set Batches loaded = {len(eval_dataloader)}")
 
         longest_seq_length, _ = get_longest_seq_length(
             torch.utils.data.ConcatDataset([train_dataloader.dataset, eval_dataloader.dataset])
@@ -329,7 +331,7 @@ def main(peft_config_file: str = None, **kwargs) -> None:
 
     # Create DataLoaders for the training and validation dataset
     train_dataloader, eval_dataloader, longest_seq_length = setup_dataloaders(train_config, dataset_config, tokenizer)
-    logger.info(
+    logger.log_rank_zero(
         f"The longest sequence length in the train data is {longest_seq_length}, "
         f"passed context length is {train_config.context_length} and overall model's context length is "
         f"{model.config.max_position_embeddings}"

QEfficient/finetune/dataset/custom_dataset.py

@@ -8,7 +8,7 @@
 import importlib
 from pathlib import Path
 
-from QEfficient.utils.logging_utils import ft_logger as logger
+from QEfficient.utils.logging_utils import logger
 
 
 def load_module_from_py_file(py_file: str) -> object:

QEfficient/finetune/utils/train_utils.py

@@ -19,7 +19,7 @@
 from tqdm import tqdm
 
 from QEfficient.finetune.configs.training import TrainConfig
-from QEfficient.utils.logging_utils import ft_logger as logger
+from QEfficient.utils.logging_utils import logger
 
 try:
     import torch_qaic  # noqa: F401
@@ -28,7 +28,7 @@
     import torch_qaic.utils as qaic_utils  # noqa: F401
     from torch.qaic.amp import GradScaler as QAicGradScaler
 except ImportError as e:
-    logger.warning(f"{e}. Moving ahead without these qaic modules.")
+    logger.log_rank_zero(f"{e}. Moving ahead without these qaic modules.")
 
 from torch.amp import GradScaler
 
@@ -111,21 +111,21 @@ def train(
     for epoch in range(train_config.num_epochs):
         if loss_0_counter.item() == train_config.convergence_counter:
             if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
-                logger.info(
+                logger.log_rank_zero(
                     f"Skipping epoch {epoch + 1} since loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps."
                 )
                 break
 
         if train_config.use_peft and train_config.from_peft_checkpoint:
             intermediate_epoch = int(train_config.from_peft_checkpoint.split("/")[-2].split("_")[-1]) - 1
             if epoch < intermediate_epoch:
-                logger.info(f"Skipping epoch {epoch + 1} since fine tuning has already completed for it.")
+                logger.log_rank_zero(f"Skipping epoch {epoch + 1} since fine tuning has already completed for it.")
                 # to bring the count of train_step in sync with where it left off
                 total_train_steps += len(train_dataloader)
                 continue
 
-        logger.info(f"Starting epoch {epoch + 1}/{train_config.num_epochs}")
-        logger.info(f"train_config.max_train_step: {train_config.max_train_step}")
+        logger.log_rank_zero(f"Starting epoch {epoch + 1}/{train_config.num_epochs}")
+        logger.log_rank_zero(f"train_config.max_train_step: {train_config.max_train_step}")
         # stop when the maximum number of training steps is reached
         if max_steps_reached:
             break
@@ -152,7 +152,7 @@ def train(
                 # to bring the count of train_step in sync with where it left off
                 if epoch == intermediate_epoch and step == 0:
                     total_train_steps += intermediate_step
-                    logger.info(
+                    logger.log_rank_zero(
                         f"Skipping first {intermediate_step} steps for epoch {epoch + 1}, since fine tuning has already completed for it."
                     )
                 if epoch == intermediate_epoch and step < intermediate_step:
@@ -266,7 +266,7 @@ def train(
                 )
             if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
                 if loss_0_counter.item() == train_config.convergence_counter:
-                    logger.info(
+                    logger.log_rank_zero(
                         f"Loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps.Hence,stopping the fine tuning."
                     )
                     break
@@ -328,15 +328,15 @@ def train(
         if train_config.run_validation:
             if eval_epoch_loss < best_val_loss:
                 best_val_loss = eval_epoch_loss
-                logger.info(f"best eval loss on epoch {epoch + 1} is {best_val_loss}")
+                logger.log_rank_zero(f"best eval loss on epoch {epoch + 1} is {best_val_loss}")
             val_loss.append(float(eval_epoch_loss))
             val_metric.append(float(eval_metric))
         if train_config.task_type == "seq_classification":
-            logger.info(
+            logger.log_rank_zero(
                 f"Epoch {epoch + 1}: train_acc={metric_val:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
             )
         else:
-            logger.info(
+            logger.log_rank_zero(
                 f"Epoch {epoch + 1}: train_metric={metric_val:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
             )
 
@@ -440,7 +440,7 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
         eval_metric = torch.exp(eval_epoch_loss)
 
     # Print evaluation metrics
-    logger.info(f"{eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
+    logger.log_rank_zero(f"{eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
 
     return eval_epoch_loss, eval_metric, val_step_loss, val_step_metric
 
@@ -469,10 +469,8 @@ def print_model_size(model, config) -> None:
         model: The PyTorch model.
         model_name (str): Name of the model.
     """
-
-    logger.info(f"Model : {config.model_name}")
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    logger.info(f"{config.model_name} has {total_params / 1e6} Million params\n")
+    logger.log_rank_zero(f"{config.model_name} has {total_params / 1e6} Million params\n")
 
 
 def save_to_json(

QEfficient/utils/logging_utils.py

@@ -6,6 +6,12 @@
 # -----------------------------------------------------------------------------
 
 import logging
+import os
+from datetime import datetime
+
+import torch.distributed as dist
+
+from QEfficient.utils.constants import ROOT_DIR
 
 
 class QEffFormatter(logging.Formatter):
@@ -38,82 +44,48 @@ def format(self, record):
         return formatter.format(record)
 
 
-def create_logger() -> logging.Logger:
+def create_logger(level=logging.INFO, dump_logs=True) -> logging.Logger:
     """
     Creates a logger object with Colored QEffFormatter.
     """
     logger = logging.getLogger("QEfficient")
 
-    # create console handler and set level to debug
+    # create console handler and set level
     ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
-    # define formatter
+    ch.setLevel(level)
     ch.setFormatter(QEffFormatter())
-
     logger.addHandler(ch)
-    return logger
-
-
-# Define the logger object that can be used for logging purposes throughout the module.
-logger = create_logger()
-
-
-def create_ft_logger(log_file="finetune.log") -> logging.Logger:
-    """
-    Creates a logger object with Colored QEffFormatter.
-    """
-    logger = logging.getLogger("QEfficient")
 
-    # create console handler and set level to debug
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
-    ch.setFormatter(QEffFormatter())
-    logger.addHandler(ch)
+    if dump_logs:
+        logs_path = os.path.join(ROOT_DIR, "logs")
+        if not os.path.exists(logs_path):
+            os.makedirs(logs_path, exist_ok=True)
+        file_name = f"log-file-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}" + ".txt"
+        log_file = os.path.join(logs_path, file_name)
 
-    # create file handler and set level to debug
-    fh = logging.FileHandler(log_file)
-    fh.setLevel(logging.INFO)
-    fh.setFormatter(QEffFormatter())
-    logger.addHandler(fh)
+        # create file handler and set level
+        fh = logging.FileHandler(log_file)
+        fh.setLevel(level)
+        formatter = logging.Formatter("%(levelname)s - %(name)s - %(message)s")
+        fh.setFormatter(formatter)
+        logger.addHandler(fh)
 
     return logger
 
 
-# Define the logger object that can be used for logging purposes throughout the finetuning module.
-ft_logger = create_ft_logger()
-"""
-
-class FT_Logger:
-    def __init__(self, level=logging.INFO, log_file="finetune.log"):
-        self.logger = logging.getLogger("QEfficient")
-        self.logger.setLevel(level)
-        self.level = level
-
-        # Create handlers
-        self.file_handler = logging.FileHandler(log_file)
-        self.console_handler = logging.StreamHandler()
-
-        self.file_handler.setFormatter(QEffFormatter())
-        self.console_handler.setFormatter(QEffFormatter())
+class CustomLogger(logging.Logger):
+    def raise_runtimeerror(self, message):
+        self.error(message)
+        raise RuntimeError(message)
 
-        # Add handlers to the logger
-        self.logger.addHandler(self.file_handler)
-        self.logger.addHandler(self.console_handler)
+    def log_rank_zero(self, msg: str, level: int = logging.INFO) -> None:
+        rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+        if rank != 0:
+            return
+        self.log(level, msg, stacklevel=2)
 
-    def get_logger(self):
-        return self.logger
-        
-    def raise_valueerror(self, message):
-        self.logger.error(message)
-        raise ValueError(message)
 
-    def raise_runtimeerror(self, message):
-        self.logger.error(message)
-        raise RuntimeError(message)
-        
-    def raise_filenotfounderror(self, message):
-        self.logger.error(message)
-        raise FileNotFoundError(message)
+logging.setLoggerClass(CustomLogger)
 
-ft_logger = FT_Logger().get_logger()
-"""
+# Define the logger object that can be used for logging purposes throughout the module.
+logger = create_logger()