query deepspeed global grad norm (bigscience-workshop#8)

stas00 · Aug 10, 2021 · 4e27f6d · 4e27f6d
1 parent 8947494
commit 4e27f6d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/training.py b/megatron/training.py
@@ -380,7 +380,7 @@ def train_step(forward_step_func, data_iterator,
         assert isinstance(model[0], deepspeed.PipelineEngine), model
         loss = model[0].train_batch(data_iter=data_iterator)
         skipped_iter = 0
-        grad_norm = 0.
+        grad_norm = model[0].get_global_grad_norm()
         num_zeros_in_grad = 0
         return {'lm loss' : loss}, skipped_iter, grad_norm, num_zeros_in_grad