bug/feature: Fixing mixed precision training

src/refiners/training_utils/trainer.py

@@ -7,6 +7,7 @@
 from loguru import logger
 from torch import Tensor, device as Device, dtype as DType, nn
 from torch.autograd import backward
+from torch.cuda.amp import GradScaler
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import (
     CosineAnnealingLR,
@@ -183,6 +184,12 @@ def dtype(self) -> DType:
         logger.info(f"Using dtype: {dtype}")
         return dtype
 
+    @cached_property
+    def scaler(self) -> GradScaler | None:
+        if self.config.training.dtype == "float32":
+            return None
+        return GradScaler()
+
     @property
     def learnable_parameters(self) -> list[nn.Parameter]:
         """Returns a list of learnable parameters in all models"""
@@ -341,15 +348,34 @@ def compute_loss(self, batch: Batch) -> Tensor:
     def compute_evaluation(self) -> None:
         pass
 
+    def backward_step(self, scaled_loss: Tensor) -> None:
+        if self.scaler is None:
+            backward(tensors=scaled_loss)
+            return
+        self.scaler.scale(scaled_loss).backward()  # type: ignore
+
+    # logic from accelerator
+    def optimizer_step(self) -> None:
+        if self.scaler is None:
+            self.optimizer.step()
+            return
+        scale_before = self.scaler.get_scale()  # type: ignore
+        self.scaler.step(self.optimizer)  # type: ignore
+        self.scaler.update()  # type: ignore
+        scale_after = self.scaler.get_scale()  # type: ignore
+        # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
+        if scale_after < scale_before:
+            logger.info("Overflow in optimizer caused optimizer to skip")
+
     def backward(self) -> None:
         """Backward pass on the loss."""
         self._call_callbacks(event_name="on_backward_begin")
         scaled_loss = self.loss / self.clock.num_step_per_iteration
-        backward(tensors=scaled_loss)
+        self.backward_step(scaled_loss)
         self._call_callbacks(event_name="on_backward_end")
         if self.clock.is_optimizer_step:
             self._call_callbacks(event_name="on_optimizer_step_begin")
-            self.optimizer.step()
+            self.optimizer_step()
             self.optimizer.zero_grad()
             self._call_callbacks(event_name="on_optimizer_step_end")
         if self.clock.is_lr_scheduler_step: