FT logger

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

Author: mamtsing

QEfficient/cloud/finetune.py

@@ -110,7 +110,7 @@ def load_model_and_tokenizer(
         - Resizes model embeddings if tokenizer vocab size exceeds model embedding size.
         - Sets pad_token_id to eos_token_id if not defined in the tokenizer.
     """
-    logger.log_rank_zero(f"loading HuggingFace model for {train_config.model_name}")
+    logger.log_rank_zero(f"Loading HuggingFace model for {train_config.model_name}")
     pretrained_model_path = hf_download(train_config.model_name)
     if train_config.task_type == "seq_classification":
         model = AutoModelForSequenceClassification.from_pretrained(
@@ -149,8 +149,7 @@ def load_model_and_tokenizer(
         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.
-    print_model_size(model, train_config)
+    print_model_size(model)
 
     # Note: Need to call this before calling PeftModel.from_pretrained or get_peft_model.
     # Because, both makes model.is_gradient_checkpointing = True which is used in peft library to
@@ -301,7 +300,7 @@ def main(peft_config_file: str = None, **kwargs) -> None:
     scheduler = StepLR(optimizer, step_size=1, gamma=train_config.gamma)
     if train_config.enable_ddp:
         model = nn.parallel.DistributedDataParallel(model, device_ids=[dist.get_rank()])
-    _ = train(
+    results = train(
         model,
         tokenizer,
         train_dataloader,
@@ -313,7 +312,7 @@ def main(peft_config_file: str = None, **kwargs) -> None:
     )
     if train_config.enable_ddp:
         dist.destroy_process_group()
-    return
+    return results
 
 
 if __name__ == "__main__":

QEfficient/finetune/dataset/grammar_dataset.py

@@ -22,7 +22,7 @@ def __init__(self, tokenizer, csv_name=None, context_length=None):
                 delimiter=",",
             )
         except Exception as e:
-            logger.error(
+            logger.raise_runtimeerror(
                 "Loading of grammar dataset failed! Please check (https://github.com/meta-llama/llama-recipes/blob/main/src/llama_recipes/datasets/grammar_dataset/grammar_dataset_process.ipynb) for details on how to download the dataset."
             )
             raise e

QEfficient/finetune/utils/dataset_utils.py

@@ -11,6 +11,7 @@
 
 from QEfficient.finetune.data.sampler import DistributedLengthBasedBatchSampler
 from QEfficient.finetune.dataset.dataset_config import DATALOADER_COLLATE_FUNC, DATASET_PREPROC
+from QEfficient.finetune.utils.logging_utils import logger
 
 
 def get_preprocessed_dataset(
@@ -72,7 +73,7 @@ def get_dataloader(tokenizer, dataset_config, train_config, split: str = "train"
         print("custom_data_collator is used")
         dl_kwargs["collate_fn"] = custom_data_collator
 
-    print(f"length of dataset_{split}", len(dataset))
+    logger.log_rank_zero(f"Length of {split} dataset is {len(dataset)}")
 
     # Create data loader
     dataloader = torch.utils.data.DataLoader(

QEfficient/finetune/utils/logging_utils.py

@@ -14,79 +14,44 @@
 from QEfficient.utils.constants import ROOT_DIR
 
 
-class QEffFormatter(logging.Formatter):
-    """
-    Formatter class used to set colors for printing different logging levels of messages on console.
-    """
-
-    cyan: str = "\x1b[38;5;14m"
-    yellow: str = "\x1b[33;20m"
-    red: str = "\x1b[31;20m"
-    bold_red: str = "\x1b[31;1m"
-    reset: str = "\x1b[0m"
-    common_format: str = "%(levelname)s - %(name)s - %(message)s"  # type: ignore
-    format_with_line_info = "%(levelname)s - %(name)s - %(message)s  (%(filename)s:%(lineno)d)"  # type: ignore
-
-    FORMATS = {
-        logging.DEBUG: cyan + format_with_line_info + reset,
-        logging.INFO: cyan + common_format + reset,
-        logging.WARNING: yellow + common_format + reset,
-        logging.ERROR: red + format_with_line_info + reset,
-        logging.CRITICAL: bold_red + format_with_line_info + reset,
-    }
-
-    def format(self, record):
-        """
-        Overriding the base class method to Choose format based on log level.
-        """
-        log_fmt = self.FORMATS.get(record.levelno)
-        formatter = logging.Formatter(log_fmt)
-        return formatter.format(record)
-
-
-def create_logger() -> logging.Logger:
-    """
-    Creates a logger object with Colored QEffFormatter.
-    """
-    logger = logging.getLogger("QEfficient")
-
-    # create console handler and set level
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
-    ch.setFormatter(QEffFormatter())
-    logger.addHandler(ch)
-
-    return logger
-
-
-class CustomLogger(logging.Logger):
-    def raise_runtimeerror(self, message):
-        self.error(message)
-        raise RuntimeError(message)
-
-    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 prepare_dump_logs(self, dump_logs=False):
-        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
-            fh = logging.FileHandler(log_file)
-            fh.setLevel(logging.INFO)
-            formatter = logging.Formatter("%(levelname)s - %(name)s - %(message)s")
-            fh.setFormatter(formatter)
-            logger.addHandler(fh)
-
-
-logging.setLoggerClass(CustomLogger)
-
-# Define the logger object that can be used for logging purposes throughout the module.
-logger = create_logger()
+class FTLogger:
+    def __init__(self, level=logging.DEBUG):
+        self.logger = logging.getLogger("QEfficient")
+        if not getattr(self.logger, "_custom_methods_added", False):
+            self._bind_custom_methods()
+            self.logger._custom_methods_added = True  # Prevent adding handlers/methods twice
+
+    def _bind_custom_methods(self):
+        def raise_runtimeerror(message):
+            self.logger.error(message)
+            raise RuntimeError(message)
+
+        def log_rank_zero(msg: str, level: int = logging.INFO):
+            rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+            if rank != 0:
+                return
+            self.logger.log(level, msg, stacklevel=2)
+
+        def prepare_dump_logs(dump_logs=False, level=logging.INFO):
+            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)
+
+                fh = logging.FileHandler(log_file)
+                fh.setLevel(level)
+                formatter = logging.Formatter("%(levelname)s - %(name)s - %(message)s")
+                fh.setFormatter(formatter)
+                self.logger.addHandler(fh)
+
+        self.logger.raise_runtimeerror = raise_runtimeerror
+        self.logger.log_rank_zero = log_rank_zero
+        self.logger.prepare_dump_logs = prepare_dump_logs
+
+    def get_logger(self):
+        return self.logger
+
+
+logger = FTLogger().get_logger()

QEfficient/finetune/utils/train_utils.py

@@ -84,7 +84,7 @@ def train(
     max_steps_reached = False  # Flag to indicate max training steps reached
 
     tensorboard_updates = None
-    if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
+    if (not train_config.enable_ddp) or (local_rank == 0):
         tensorboard_updates = SummaryWriter()
 
     device_type = torch.device(device).type
@@ -215,7 +215,7 @@ def train(
                 else:
                     loss_0_counter = torch.tensor([0]).to(device)
 
-            if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
+            if (not train_config.enable_ddp) or (local_rank == 0):
                 tensorboard_updates.add_scalars("loss", {"train": loss}, total_train_steps)
 
             if train_config.save_metrics:
@@ -300,18 +300,10 @@ def train(
         lr_scheduler.step()
 
         if train_config.run_validation:
-            if train_config.enable_ddp:
-                dist.barrier()
-                eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation_helper(
-                    model, train_config, eval_dataloader, device
-                )
-                if local_rank == 0:
-                    tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
-
-            else:
-                eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation_helper(
-                    model, train_config, eval_dataloader, device
-                )
+            eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation_helper(
+                model, train_config, eval_dataloader, device
+            )
+            if (not train_config.enable_ddp) or (local_rank == 0):
                 tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
 
             if train_config.save_metrics:
@@ -385,6 +377,9 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
 
     Returns: eval_epoch_loss, eval_metric, eval_step_loss, eval_step_metric
     """
+    if train_config.enable_ddp:
+        dist.barrier()
+
     model.eval()
 
     if train_config.task_type == "seq_classification":
@@ -457,16 +452,15 @@ def get_longest_seq_length(data: List[Dict]) -> Tuple[int, int]:
     return longest_seq_length, longest_seq_ix
 
 
-def print_model_size(model, config) -> None:
+def print_model_size(model) -> None:
     """
     Print model name, the number of trainable parameters and initialization time.
 
     Args:
-        model: The PyTorch model.
-        config : Config of the model.
+        model: PyTorch model.
     """
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    logger.log_rank_zero(f"{config.model_name} has {total_params / 1e6} Million params.")
+    logger.log_rank_zero(f"Model has {total_params / 1e6} Million params.")
 
 
 def print_trainable_parameters(model) -> None:
@@ -478,7 +472,7 @@ def print_trainable_parameters(model) -> None:
     """
     trainable_params, all_param = model.get_nb_trainable_parameters()
     logger.log_rank_zero(
-        f"trainable params: {trainable_params:,d} || all params: {all_param:,d} || trainable%: {100 * trainable_params / all_param:.4f}"
+        f"Trainable params: {trainable_params:,d} || all params: {all_param:,d} || trainable%: {100 * trainable_params / all_param:.4f}"
     )