adding gradient norm tracking when grad clipping is enabled

Author: azrael417

makani/utils/training/autoencoder_trainer.py

@@ -491,6 +491,11 @@ def train_one_epoch(self, profiler=None):
         # we need this for the loss average
         accumulated_loss = torch.zeros((2), dtype=torch.float32, device=self.device)
 
+        if self.max_grad_norm > 0.0:
+            accumulated_grad_norm = torch.zeros((2), dtype=torch.float32, device=self.device, requires_grad=False)
+        else:
+            accumulated_grad_norm = None
+
         train_steps = 0
         train_start = time.perf_counter_ns()
         self.model_train.zero_grad(set_to_none=True)
@@ -535,7 +540,9 @@ def train_one_epoch(self, profiler=None):
             if do_update:
                 if self.max_grad_norm > 0.0:
                     self.gscaler.unscale_(self.optimizer)
-                    clip_grads(self.model_train, self.max_grad_norm)
+                    grad_norm = clip_grads(self.model_train, self.max_grad_norm)
+                    accumulated_grad_norm[0] += grad_norm.detach()
+                    accumulated_grad_norm[1] += 1.0
 
                 self.gscaler.step(self.optimizer)
                 self.gscaler.update()
@@ -570,6 +577,11 @@ def train_one_epoch(self, profiler=None):
         # add train steps to log
         logs["train_steps"] = train_steps
 
+        # log gradient norm
+        if accumulated_grad_norm is not None:
+            grad_norm = accumulated_grad_norm[0] / accumulated_grad_norm[1]
+            logs["gradient norm"] = grad_norm.item()
+
         # global sync is in order
         if dist.is_initialized():
             dist.barrier(device_ids=[self.device.index])
@@ -719,6 +731,9 @@ def get_pad(nchar):
             # validation summary
             self.logger.info("Metrics:")
             self.logger.info(print_prefix + "training loss: {}{}".format(get_pad(pad_len[0]), train_logs["loss"]))
+            if "gradient norm" in train_logs:
+                plen = max_len - len("gradient norm")
+                self.logger.info(print_prefix + "gradient norm: {}{}".format(get_pad(plen), train_logs["gradient norm"]))
             self.logger.info(print_prefix + "validation loss: {}{}".format(get_pad(pad_len[1]), valid_logs["base"]["validation loss"]))
             for idk, key in enumerate(print_list[3:], start=3):
                 value = valid_logs["metrics"][key]

makani/utils/training/deterministic_trainer.py

@@ -490,6 +490,11 @@ def train_one_epoch(self, profiler=None):
         # we need this for the loss average
         accumulated_loss = torch.zeros((2), dtype=torch.float32, device=self.device, requires_grad=False)
 
+        if self.max_grad_norm > 0.0:
+            accumulated_grad_norm = torch.zeros((2), dtype=torch.float32, device=self.device, requires_grad=False)
+        else:
+            accumulated_grad_norm = None
+
         train_steps = 0
         train_start = time.perf_counter_ns()
         self.model_train.zero_grad(set_to_none=True)
@@ -543,7 +548,9 @@ def train_one_epoch(self, profiler=None):
             if do_update:
                 if self.max_grad_norm > 0.0:
                     self.gscaler.unscale_(self.optimizer)
-                    clip_grads(self.model_train, self.max_grad_norm)
+                    grad_norm = clip_grads(self.model_train, self.max_grad_norm)
+                    accumulated_grad_norm[0] += grad_norm.detach()
+                    accumulated_grad_norm[1] += 1.0
 
                 self.gscaler.step(self.optimizer)
                 self.gscaler.update()
@@ -581,6 +588,11 @@ def train_one_epoch(self, profiler=None):
         # add train steps to log
         logs["train_steps"] = train_steps
 
+        # log gradient norm
+        if accumulated_grad_norm is not None:
+            grad_norm = accumulated_grad_norm[0] / accumulated_grad_norm[1]
+            logs["gradient norm"] = grad_norm.item()
+
         # global sync is in order
         if dist.is_initialized():
             dist.barrier(device_ids=[self.device.index])
@@ -725,6 +737,9 @@ def get_pad(nchar):
             # validation summary
             self.logger.info("Metrics:")
             self.logger.info(print_prefix + "training loss: {}{}".format(get_pad(pad_len[0]), train_logs["loss"]))
+            if "gradient norm" in train_logs:
+                plen = max_len - len("gradient norm")
+                self.logger.info(print_prefix + "gradient norm: {}{}".format(get_pad(plen), train_logs["gradient norm"]))
             self.logger.info(print_prefix + "validation loss: {}{}".format(get_pad(pad_len[1]), valid_logs["base"]["validation loss"]))
             for idk, key in enumerate(print_list[3:], start=3):
                 value = valid_logs["metrics"][key]

makani/utils/training/ensemble_trainer.py

@@ -503,6 +503,11 @@ def train_one_epoch(self, profiler=None):
         # we need this for the loss average
         accumulated_loss = torch.zeros((2), dtype=torch.float32, device=self.device)
 
+        if self.max_grad_norm > 0.0:
+            accumulated_grad_norm = torch.zeros((2), dtype=torch.float32, device=self.device)
+        else:
+            accumulated_grad_norm = None
+
         train_steps = 0
         train_start = time.perf_counter_ns()
         self.model_train.zero_grad(set_to_none=True)
@@ -552,7 +557,9 @@ def train_one_epoch(self, profiler=None):
             if do_update:
                 if self.max_grad_norm > 0.0:
                     self.gscaler.unscale_(self.model_optimizer)
-                    clip_grads(self.model_train, self.max_grad_norm)
+                    grad_norm = clip_grads(self.model_train, self.max_grad_norm)
+                    accumulated_grad_norm[0] += grad_norm.detach()
+                    accumulated_grad_norm[1] += 1.0
 
                 self.gscaler.step(self.model_optimizer)
                 self.gscaler.update()
@@ -590,6 +597,11 @@ def train_one_epoch(self, profiler=None):
         # add train steps to log
         logs["train_steps"] = train_steps
 
+        # log gradient norm
+        if accumulated_grad_norm is not None:
+            grad_norm = accumulated_grad_norm[0] / accumulated_grad_norm[1]
+            logs["gradient norm"] = grad_norm.item()
+
         # global sync is in order
         if dist.is_initialized():
             dist.barrier(device_ids=[self.device.index])
@@ -774,6 +786,9 @@ def get_pad(nchar):
             # validation summary
             self.logger.info("Metrics:")
             self.logger.info(print_prefix + "training loss: {}{}".format(get_pad(pad_len[0]), train_logs["loss"]))
+            if "gradient norm" in train_logs:
+                plen = max_len - len("gradient norm")
+                self.logger.info(print_prefix + "gradient norm: {}{}".format(get_pad(plen), train_logs["gradient norm"]))
             self.logger.info(print_prefix + "validation loss: {}{}".format(get_pad(pad_len[1]), valid_logs["base"]["validation loss"]))
             for idk, key in enumerate(print_list[3:], start=3):
                 value = valid_logs["metrics"][key]

makani/utils/training/stochastic_trainer.py

@@ -472,6 +472,11 @@ def train_one_epoch(self):
         # we need this for the loss average
         accumulated_loss = torch.zeros((2), dtype=torch.float32, device=self.device)
 
+        if self.max_grad_norm > 0.0:
+            accumulated_grad_norm = torch.zeros((2), dtype=torch.float32, device=self.device, requires_grad=False)
+        else:
+            accumulated_grad_norm = None
+
         train_steps = 0
         train_start = time.perf_counter_ns()
         self.model_train.zero_grad(set_to_none=True)
@@ -518,7 +523,9 @@ def train_one_epoch(self):
             if do_update:
                 if self.max_grad_norm > 0.0:
                     self.gscaler.unscale_(self.optimizer)
-                    clip_grads(self.model_train, self.max_grad_norm)
+                    grad_norm = clip_grads(self.model_train, self.max_grad_norm)
+                    accumulated_grad_norm[0] += grad_norm.detach()
+                    accumulated_grad_norm[1] += 1.0
 
                 # perform weight update
                 self.gscaler.step(self.optimizer)
@@ -556,6 +563,11 @@ def train_one_epoch(self):
         # add train steps to log
         logs["train_steps"] = train_steps
 
+        # log gradient norm
+        if accumulated_grad_norm is not None:
+            grad_norm = accumulated_grad_norm[0] / accumulated_grad_norm[1]
+            logs["gradient norm"] = grad_norm.item()
+
         # global sync is in order
         if dist.is_initialized():
             dist.barrier(device_ids=[self.device.index])
@@ -710,6 +722,9 @@ def get_pad(nchar):
             # validation summary
             self.logger.info("Metrics:")
             self.logger.info(print_prefix + "training loss: {}{}".format(get_pad(pad_len[0]), train_logs["loss"]))
+            if "gradient norm" in train_logs:
+                plen = max_len - len("gradient norm")
+                self.logger.info(print_prefix + "gradient norm: {}{}".format(get_pad(plen), train_logs["gradient norm"]))
             self.logger.info(print_prefix + "validation loss: {}{}".format(get_pad(pad_len[1]), valid_logs["base"]["validation loss"]))
             for idk, key in enumerate(print_list[3:], start=3):
                 value = valid_logs["metrics"][key]

makani/utils/training/training_helpers.py

@@ -99,7 +99,7 @@ def clip_grads(model, max_grad_norm, norm_type=2.0):
 
             param.grad.mul_(clip_factor)
 
-    return
+    return total_gnorm
 
 
 def wandb_register_activations_monitor(model: nn.Module, step: int):