From 22de429cef7f41cd4cf643945262a2c8eb6e16c7 Mon Sep 17 00:00:00 2001
From: Loubna Ben Allal <44069155+loubnabnl@users.noreply.github.com>
Date: Fri, 26 May 2023 10:09:45 +0200
Subject: [PATCH] Add tokens-per-second-per-gpu to the printed logs instead of
 just wandb logs

---
 megatron/training.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/megatron/training.py b/megatron/training.py
index 468a600291b..c02a91f39fa 100644
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -655,6 +655,7 @@ def add_to_logging(name):
         log_string += ' number of nan iterations: {:3d} |'.format(
             total_loss_dict[nan_iters_key])
         log_string += ' TFLOPs: {:.2f} |'.format(tflops)
+        log_string += ' tokens-per-second-per-gpu: {:.2f} |'.format(tokens_per_sec_per_gpu)
         total_loss_dict[advanced_iters_key] = 0
         total_loss_dict[skipped_iters_key] = 0
         total_loss_dict[nan_iters_key] = 0