Moved eval all reduce code after including it in json list.

Signed-off-by: meetkuma <[email protected]>

Author: quic-meetkuma

QEfficient/finetune/utils/train_utils.py

@@ -369,13 +369,9 @@ def train(
             eval_loss, eval_metric, step_loss, step_metric = evaluation_helper(
                 model, train_config, eval_dataloader, device
             )
-            # Print evaluation metrics
-            print(
-                f"Epoch {epoch + 1}: Eval Loss: {eval_loss.detach().cpu():.4f}, Eval metric: {eval_metric.detach().cpu():.4f}"
-            )
             if eval_loss < best_val_loss:
                 best_val_loss = eval_loss
-                print(f"best eval loss on epoch {epoch + 1} is {best_val_loss:.4f}")
+                print(f"Best eval loss on epoch {epoch + 1} is {best_val_loss:.4f}")
 
             if is_rank_zero():
                 tensorboard_updates.add_scalars("loss", {"eval": eval_loss}, total_train_steps)
@@ -385,6 +381,16 @@ def train(
                 val_loss.append(float(eval_loss))
                 val_metric.append(float(eval_metric))
 
+            if train_config.enable_ddp:
+                dist.all_reduce(eval_loss, op=dist.ReduceOp.SUM)
+                eval_loss /= get_num_ddp_devices()
+                dist.all_reduce(eval_metric, op=dist.ReduceOp.SUM)
+                eval_metric /= get_num_ddp_devices()
+
+            print(
+                f"Epoch {epoch + 1}: Eval Loss: {eval_loss.detach().cpu():.4f}, Eval metric: {eval_metric.detach().cpu():.4f}"
+            )
+
         # saving the adapters after completion of each epoch
         if train_config.save_model:
             if train_config.enable_ddp:
@@ -507,12 +513,6 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
     else:
         eval_metric = torch.exp(eval_loss)
 
-    if train_config.enable_ddp:
-        dist.all_reduce(eval_loss, op=dist.ReduceOp.SUM)
-        eval_loss /= get_num_ddp_devices()
-        dist.all_reduce(eval_metric, op=dist.ReduceOp.SUM)
-        eval_metric /= get_num_ddp_devices()
-
     return eval_loss, eval_metric, val_step_loss, val_step_metric