Address comments

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

Author: quic-mamta

QEfficient/cloud/finetune.py

@@ -5,6 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
+import logging
 import random
 import warnings
 from typing import Any, Dict, Optional, Union
@@ -18,7 +19,7 @@
 import torch.utils.data
 from peft import PeftModel, get_peft_model
 from torch.optim.lr_scheduler import StepLR
-from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModel, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
 
 from QEfficient.finetune.configs.training import TrainConfig
 from QEfficient.finetune.utils.config_utils import (
@@ -41,8 +42,8 @@
 except ImportError as e:
     logger.warning(f"{e}. Moving ahead without these qaic modules.")
 
+logger.setLevel(logging.INFO)
 
-from transformers import AutoModelForSequenceClassification
 
 # Suppress all warnings
 warnings.filterwarnings("ignore")
@@ -245,7 +246,7 @@ def setup_dataloaders(
     #         )
     ##
     train_dl_kwargs = get_dataloader_kwargs(train_config, dataset_train, dataset_processer, "train")
-    logger.info("length of dataset_train", len(dataset_train))
+    logger.info(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)
@@ -260,7 +261,7 @@ def setup_dataloaders(
         pin_memory=True,
         **train_dl_kwargs,
     )
-    logger.info(f"--> Num of Training Set Batches loaded = {len(train_dataloader)}")
+    logger.info(f"Num of Training Set Batches loaded = {len(train_dataloader)}")
 
     eval_dataloader = None
     if train_config.run_validation:
@@ -284,7 +285,7 @@ def setup_dataloaders(
                 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.info(f"Num 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])

QEfficient/finetune/dataset/grammar_dataset.py

@@ -23,7 +23,7 @@ def __init__(self, tokenizer, csv_name=None, context_length=None):
             )
         except Exception as e:
             logger.error(
-                "Loading of grammar dataset failed! Please see [here](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."
+                "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/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.info(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.info(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/config_utils.py

@@ -54,7 +54,7 @@ def update_config(config, **kwargs):
                         raise ValueError(f"Config '{config_name}' does not have parameter: '{param_name}'")
             else:
                 config_type = type(config).__name__
-                logger.warning(f"Unknown parameter '{k}' for config type '{config_type}'")
+                logger.debug(f"Unknown parameter '{k}' for config type '{config_type}'")
 
 
 def generate_peft_config(train_config: TrainConfig, peft_config_file: str = None, **kwargs) -> Any:

QEfficient/finetune/utils/train_utils.py

@@ -85,10 +85,7 @@ def train(
     device_type = device.split(":")[0]
 
     tensorboard_updates = None
-    if train_config.enable_ddp:
-        if local_rank == 0:
-            tensorboard_updates = SummaryWriter()
-    else:
+    if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
         tensorboard_updates = SummaryWriter()
 
     if train_config.grad_scaler:
@@ -113,14 +110,9 @@ def train(
     # Start the training loop
     for epoch in range(train_config.num_epochs):
         if loss_0_counter.item() == train_config.convergence_counter:
-            if train_config.enable_ddp:
-                logger.info(
-                    f"Not proceeding with epoch {epoch + 1} on device {local_rank} since loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps."
-                )
-                break
-            else:
+            if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
                 logger.info(
-                    f"Not proceeding with epoch {epoch + 1} since loss value has been <= {train_config.convergence_loss}  for last {loss_0_counter.item()} steps."
+                    f"Skipping epoch {epoch + 1} since loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps."
                 )
                 break
 
@@ -161,7 +153,7 @@ def train(
                 if epoch == intermediate_epoch and step == 0:
                     total_train_steps += intermediate_step
                     logger.info(
-                        f"skipping first {intermediate_step} steps for epoch {epoch + 1}, since fine tuning has already completed for them."
+                        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:
                     total_train_steps += 1
@@ -221,10 +213,7 @@ def train(
                 else:
                     loss_0_counter = torch.tensor([0]).to(device)
 
-            if train_config.enable_ddp:
-                if local_rank == 0:
-                    tensorboard_updates.add_scalars("loss", {"train": loss}, total_train_steps)
-            else:
+            if (not train_config.enable_ddp) or (train_config.enable_ddp and local_rank == 0):
                 tensorboard_updates.add_scalars("loss", {"train": loss}, total_train_steps)
 
             if train_config.save_metrics:
@@ -275,16 +264,10 @@ def train(
                     val_step_metric,
                     val_metric,
                 )
-            if train_config.enable_ddp:
-                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 on device {local_rank}."
-                    )
-                    break
-            else:
+            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."
+                        f"Loss value has been <= {train_config.convergence_loss} for last {loss_0_counter.item()} steps.Hence,stopping the fine tuning."
                     )
                     break
 
@@ -457,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.info(f"{eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
 
     return eval_epoch_loss, eval_metric, val_step_loss, val_step_metric
 
@@ -487,9 +470,9 @@ def print_model_size(model, config) -> None:
         model_name (str): Name of the model.
     """
 
-    logger.info(f"--> Model {config.model_name}")
+    logger.info(f"Model : {config.model_name}")
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    logger.info(f"\n--> {config.model_name} has {total_params / 1e6} Million params\n")
+    logger.info(f"{config.model_name} has {total_params / 1e6} Million params\n")
 
 
 def save_to_json(