add gradient norm tensorboard visualizatoin for deepspeed mode

2025-01-14 18:13:11 +08:00
parent a35d7bbf0e
commit 89cae4def0
1 changed files with 4 additions and 0 deletions
@@ -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: