Added per epoch time in the log files

Signed-off-by: Anusha Bhamidipati <abhamidi@qti.qualcomm.com>

Author: Anusha Bhamidipati

QEfficient/finetune/experimental/core/callbacks.py

@@ -11,6 +11,7 @@
 import time
 from typing import Any, Dict, Optional
 
+from logger import Logger
 from transformers import (
     DefaultFlowCallback,
     EarlyStoppingCallback,
@@ -33,26 +34,24 @@
 registry.callback("printer")(PrinterCallback)
 registry.callback("default_flow")(DefaultFlowCallback)
 registry.callback("tensorboard")(TensorBoardCallback)
+logger = Logger("epoch_logger", log_file="logs/training_log.log", level=logging.DEBUG)
+
 
 @registry.callback("epoch_timer")
 # Custom callback to reset epoch start time
 class EpochTimerCallback(TrainerCallback):
     def on_epoch_begin(self, args, state, control, **kwargs):
         # Record the start time of the current epoch
         self.epoch_start_time = time.time()
-        #print(f"Epoch {state.epoch} started at {time.ctime(self.epoch_start_time)}")
 
     def on_epoch_end(self, args, state, control, **kwargs):
         # Compute time elapsed for the epoch
         elapsed = time.time() - self.epoch_start_time
-        #print(f"Time taken to execute Epoch {state.epoch}: {elapsed:.2f} seconds")
-        if state.is_world_process_zero:
-              print(f"[Epoch {state.epoch:.2f}] {elapsed:.2f} sec")
-	# attach to log history so it goes to TB/W&B/etc.
+        logger.log_rank_zero(f"[Epoch {state.epoch:.2f}] {elapsed:.2f} sec")
+        # attach to log history so it goes to TB/W&B/etc.
         state.log_history.append({"train/epoch_time_sec": elapsed, "epoch": state.epoch})
         control.should_log = True
 
-
 @registry.callback("enhanced_progressbar")
 class EnhancedProgressCallback(ProgressCallback):
     """
@@ -244,10 +243,10 @@ def replace_progress_callback(trainer: Any, callbacks: list[Any], logger: Any =
             pass
 
         try:
-            #Add Epoch Timer
+            # Add Epoch Timer
             epoch_timer = ComponentFactory.create_callback("epoch_timer")
             trainer.add_callback(epoch_timer)
-            #Add EnhancedProgressCallback
+            # Add EnhancedProgressCallback
             enhanced_callback = ComponentFactory.create_callback("enhanced_progressbar")
             trainer.add_callback(enhanced_callback)
         except Exception as e:

QEfficient/finetune/experimental/core/trainer/base_trainer.py

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # -----------------------------------------------------------------------------
+
 from typing import Optional
 
 from peft import get_peft_model
@@ -62,7 +63,7 @@ def __init__(
         if peft_config is not None and model is not None:
             model = get_peft_model(model, peft_config)
             model.print_trainable_parameters()
-
+            model.config.use_cache = False
         # Initialize the parent Trainer class
         super().__init__(
             model=model,