From 89cae4def04d463b2d6518e6634c70bdcc07d10c Mon Sep 17 00:00:00 2001
From: "mengli.cml" <mengli.cml@alibaba-inc.com>
Date: Tue, 14 Jan 2025 18:13:11 +0800
Subject: [PATCH] add gradient norm tensorboard visualizatoin for deepspeed
 mode

---
 scripts/train.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/scripts/train.py b/scripts/train.py
index 10b2ca7..ab469ac 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -1989,6 +1989,10 @@ def main():
                 lr_scheduler.step()
                 optimizer.zero_grad()
 
+                if args.use_deepspeed and hasattr(optimizer, 'optimizer') and hasattr(optimizer.optimizer, '_global_grad_norm') and accelerator.is_main_process:
+                    writer.add_scalar(f'gradients/norm_sum', optimizer.optimizer._global_grad_norm,
+                            global_step=global_step)
+
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients: