Gradscaler flags

speechbrain/core.py

@@ -70,6 +70,10 @@
     "compile_using_fullgraph": False,
     "compile_using_dynamic_shape_tracing": False,
     "precision": "fp32",
+    "gradscaler_init_scale": 65536.0,
+    "gradscaler_growth_factor": 2.0,
+    "gradscaler_backoff_factor": 0.5,
+    "gradscaler_growth_interval": 2000,
     "auto_mix_prec": False,
     "bfloat16_mix_prec": False,
     "max_grad_norm": 5.0,
@@ -360,6 +364,29 @@ def parse_arguments(arg_list=None):
         help="This flag enables training with automatic mixed-precision."
         "It can be set to `fp32`, `fp16`, or `bf16`.",
     )
+    parser.add_argument(
+        "--gradscaler_init_scale",
+        type=float,
+        help="GradScaler initial scale factor.",
+    )
+    parser.add_argument(
+        "--gradscaler_growth_factor",
+        type=float,
+        help="GradScaler factor by which the scale is multiplied during "
+        "`update` if no inf/NaN gradients occur for ``growth_interval`` consecutive iterations.",
+    )
+    parser.add_argument(
+        "--gradscaler_backoff_factor",
+        type=float,
+        help="GradScaler factor by which the scale is multiplied during `update`"
+        "if inf/NaN gradients occur in an iteration.",
+    )
+    parser.add_argument(
+        "--gradscaler_growth_interval",
+        type=float,
+        help="Gradscaler number of consecutive iterations without inf/NaN gradients that must occur for the scale"
+        "to be multiplied by `growth_factor`.",
+    )
     parser.add_argument(
         "--auto_mix_prec",
         default=None,
@@ -555,6 +582,14 @@ class and override any methods for which the default behavior does not
             The location for performing computations.
         precision (str)
             One of ``fp32``, ``fp16``, ``bf16``.
+        gradscaler_init_scale (float)
+            Initial scale for the GradScaler. Default: ``65536.0``.
+        gradscaler_growth_factor (float)
+            Growth factor for the GradScaler. Default: ``2.0``.
+        gradscaler_backoff_factor (float)
+            Backoff factor for the GradScaler. Default: ``0.5``.
+        gradscaler_growth_interval (int)
+            Growth interval for the GradScaler. Default: ``2000``.
         auto_mix_prec (bool)
             If ``True``, automatic mixed-precision (fp16) is used.
             Activate it only with cuda. Note: this is a
@@ -753,7 +788,13 @@ def __init__(  # noqa: C901
         logger.info(
             f"Gradscaler enabled: {gradscaler_enabled}. Using precision: {self.precision}."
         )
-        self.scaler = torch.cuda.amp.GradScaler(enabled=gradscaler_enabled)
+        self.scaler = torch.cuda.amp.GradScaler(
+            init_scale=self.gradscaler_init_scale,
+            growth_factor=self.gradscaler_growth_factor,
+            backoff_factor=self.gradscaler_backoff_factor,
+            growth_interval=self.gradscaler_growth_interval,
+            enabled=gradscaler_enabled,
+        )
 
         self.use_amp = False
         if self.device == "cpu" and self.precision == "bf16":
@@ -1133,6 +1174,7 @@ def fit_batch(self, batch):
             scaled_loss = self.scaler.scale(
                 loss / self.grad_accumulation_factor
             )
+            self.check_loss_isfinite(scaled_loss)
             scaled_loss.backward()
 
         if should_step:
@@ -1141,6 +1183,43 @@ def fit_batch(self, batch):
         self.on_fit_batch_end(batch, outputs, loss, should_step)
         return loss.detach().cpu()
 
+    def check_loss_isfinite(self, loss):
+        """Check if the loss is finite.
+
+        If the loss is not finite, log a helpful message and increment the `nonfinite_count`.
+        If the `nonfinite_count` exceeds the `--nonfinite_patience` threshold, stop the training
+        and raise an error.
+
+        This check is particularly useful when the loss becomes NaN or inf, while the
+        parameters and gradients remain finite. It helps prevent getting stuck in an
+        infinite loop during training.
+
+        Arguments
+        ---------
+        loss : tensor
+            The loss tensor after ``backward()`` has been called but
+            before the optimizers ``step()``.
+        """
+        if not torch.isfinite(loss):
+            self.nonfinite_count += 1
+
+            # Print helpful debug info
+            logger.warning(f"Loss is {loss}.")
+            for p in self.modules.parameters():
+                if not torch.isfinite(p).all():
+                    logger.warning("Parameter is not finite: " + str(p))
+
+            # Check if patience is exhausted
+            if self.nonfinite_count > self.nonfinite_patience:
+                raise ValueError(
+                    "Loss is not finite and patience is exhausted. "
+                    "To debug, wrap `fit()` with "
+                    "autograd's `detect_anomaly()`, e.g.\n\nwith "
+                    "torch.autograd.detect_anomaly():\n\tbrain.fit(...)"
+                )
+            else:
+                logger.warning("Patience not yet exhausted.")
+
     def check_gradients(self):
         """ Checks if the gradients are finite. If not, it will emit a warning and set them to zero."""
         for param in self.modules.parameters():