Update logging_utils and log for zero rank

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

Author: mamtsing

QEfficient/cloud/finetune.py

@@ -32,17 +32,22 @@
     get_custom_data_collator,
     get_preprocessed_dataset,
 )
-from QEfficient.finetune.utils.train_utils import get_longest_seq_length, print_model_size, train
+from QEfficient.finetune.utils.train_utils import (
+    get_longest_seq_length,
+    print_model_size,
+    print_trainable_parameters,
+    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
+
+logger.setLevel(logging.INFO)
 
 # 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.setLevel(logging.INFO)
+    logger.log_rank_zero(f"{e}. Moving ahead without these qaic modules.")
 
 
 # Suppress all warnings
@@ -121,7 +126,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 +151,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 +167,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)
 
@@ -197,7 +204,7 @@ def apply_peft(
     else:
         peft_config = generate_peft_config(train_config, peft_config_file, **kwargs)
         model = get_peft_model(model, peft_config)
-    model.print_trainable_parameters()
+    print_trainable_parameters(model)
 
     return model
 
@@ -222,7 +229,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 +253,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 +268,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 +288,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 +336,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}"
@@ -340,7 +347,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()])
-    results = train(
+    _ = train(
         model,
         tokenizer,
         train_dataloader,
@@ -352,7 +359,7 @@ def main(peft_config_file: str = None, **kwargs) -> None:
     )
     if train_config.enable_ddp:
         dist.destroy_process_group()
-    return results
+    return
 
 
 if __name__ == "__main__":

QEfficient/finetune/configs/training.py

@@ -105,3 +105,5 @@ class TrainConfig:
     grad_scaler: bool = True
     dump_root_dir: str = "meta-llama-samsum-mismatches/step_"
     opByOpVerifier: bool = False
+
+    dump_logs: bool = True

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:
@@ -32,20 +32,19 @@ def get_custom_dataset(dataset_config, tokenizer, split: str):
         module_path, func_name = dataset_config.file, "get_custom_dataset"
 
     if not module_path.endswith(".py"):
-        raise ValueError(f"Dataset file {module_path} is not a .py file.")
+        logger.raise_runtimeerror(f"Dataset file {module_path} is not a .py file.")
 
     module_path = Path(module_path)
     if not module_path.is_file():
-        raise FileNotFoundError(f"Dataset py file {module_path.as_posix()} does not exist or is not a file.")
+        logger.raise_runtimeerror(f"Dataset py file {module_path.as_posix()} does not exist or is not a file.")
 
     module = load_module_from_py_file(module_path.as_posix())
     try:
         return getattr(module, func_name)(dataset_config, tokenizer, split)
-    except AttributeError as e:
-        logger.error(
+    except AttributeError:
+        logger.raise_runtimeerror(
             f"It seems like the given method name ({func_name}) is not present in the dataset .py file ({module_path.as_posix()})."
         )
-        raise e
 
 
 def get_data_collator(dataset_processer, dataset_config):
@@ -55,11 +54,11 @@ def get_data_collator(dataset_processer, dataset_config):
         module_path, func_name = dataset_config.file, "get_data_collator"
 
     if not module_path.endswith(".py"):
-        raise ValueError(f"Dataset file {module_path} is not a .py file.")
+        logger.raise_runtimeerror(f"Dataset file {module_path} is not a .py file.")
 
     module_path = Path(module_path)
     if not module_path.is_file():
-        raise FileNotFoundError(f"Dataset py file {module_path.as_posix()} does not exist or is not a file.")
+        logger.raise_runtimeerror(f"Dataset py file {module_path.as_posix()} does not exist or is not a file.")
 
     module = load_module_from_py_file(module_path.as_posix())
     try:

QEfficient/finetune/eval.py

@@ -109,13 +109,13 @@ def main(**kwargs):
             pin_memory=True,
             **val_dl_kwargs,
         )
-        logger.info(f"Num of Validation Set Batches loaded = {len(eval_dataloader)}")
+        logger.log_rank_zero(f"Num of Validation Set Batches loaded = {len(eval_dataloader)}")
         if len(eval_dataloader) == 0:
             raise ValueError(
                 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"Num of Validation Set Batches loaded = {len(eval_dataloader)}")
 
     model.to(device)
     _ = evaluation(model, train_config, eval_dataloader, None, tokenizer, device)

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
 
@@ -110,22 +110,21 @@ def train(
     # Start the training loop
     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(
-                    f"Skipping epoch {epoch + 1} since loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps."
-                )
-                break
+            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 +151,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:
@@ -264,12 +263,11 @@ def train(
                     val_step_metric,
                     val_metric,
                 )
-            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(
-                        f"Loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps.Hence,stopping the fine tuning."
-                    )
-                    break
+            if loss_0_counter.item() == train_config.convergence_counter:
+                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
 
         pbar.close()
         epoch_end_time = time.perf_counter() - epoch_start_time
@@ -328,15 +326,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 +438,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
 
@@ -467,12 +465,23 @@ def print_model_size(model, config) -> None:
 
     Args:
         model: The PyTorch model.
-        model_name (str): Name of the model.
+        config : Config 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.")
+
+
+def print_trainable_parameters(model) -> None:
+    """
+    Print the number of trainable parameters, all params and percentage of trainablke params.
+
+    Args:
+        model: The PyTorch model.
+    """
+    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}"
+    )
 
 
 def save_to_json(

QEfficient/utils/_utils.py

@@ -36,7 +36,7 @@ class DownloadRetryLimitExceeded(Exception):
 
 
 def login_and_download_hf_lm(model_name, *args, **kwargs):
-    logger.info(f"loading HuggingFace model for {model_name}")
+    logger.log_rank_zero(f"loading HuggingFace model for {model_name}")
     hf_token = kwargs.pop("hf_token", None)
     cache_dir = kwargs.pop("cache_dir", None)
     if hf_token is not None: