Log step_time

Signed-off-by: Jared Wilber <jwilber@nvidia.com>

Author: jwilber

recipes/geneformer_native_te_nvfsdp_fp8/train.py

@@ -46,6 +46,7 @@
 
 import logging
 import os
+import time
 from dataclasses import dataclass, field
 
 import hydra
@@ -217,7 +218,7 @@ def main(cfg: DictConfig) -> None:
             logger=logger,
             start_step=start_step,
         )
-
+    previous_step_time = time.perf_counter()
     for step in range(start_step, cfg.training.num_train_steps):
         # Get batch
         batch = next(dataloader)
@@ -252,6 +253,9 @@ def main(cfg: DictConfig) -> None:
 
         # Log metrics to wandb on main process
         if dist_config.is_main_process():
+            current_time = time.perf_counter()
+            step_time = current_time - previous_step_time
+            previous_step_time = current_time
             logger.info(
                 f"Step {step} loss: {loss.item()}, grad_norm: {total_norm}, lr: {optimizer.param_groups[0]['lr']}"
             )
@@ -262,6 +266,7 @@ def main(cfg: DictConfig) -> None:
                     "train/learning_rate": optimizer.param_groups[0]["lr"],
                     "train/grad_norm": total_norm,
                     "train/epoch": step / dataloader_length,
+                    "train/step_time": step_time,
                 }
             )