Track grad_norm in megatron/training.py

argonne-lcf · Jul 20, 2024 · 54bd608 · 54bd608
1 parent 019dc3c
commit 54bd608
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/training.py b/megatron/training.py
@@ -941,7 +941,7 @@ def train_step(
     # Update learning rate.
     if args.deepspeed:
         skipped_iter = 0
-        grad_norm = None
+        grad_norm = model[0].get_global_grad_norm()
         num_zeros_in_grad = None
         loss_reduced = {}
         for key in losses_reduced[0]: