bug/feature: Fixing mixed precision training

src/refiners/training_utils/config.py

@@ -22,6 +22,9 @@
 
 class TrainingConfig(BaseModel):
     device: str = "cpu"
+    automatic_mixed_precision: bool = (
+        True  # Enables automatic mixed precision which allows float32 gradients while working with lower precision.
+    )
     dtype: str = "float32"
     duration: TimeValue = {"number": 1, "unit": TimeUnit.ITERATION}
     seed: int = 0

src/refiners/training_utils/trainer.py

@@ -5,8 +5,9 @@
 
 import torch
 from loguru import logger
-from torch import Tensor, device as Device, dtype as DType, nn
+from torch import Tensor, device as Device, dtype as DType, float16, float32, nn
 from torch.autograd import backward
+from torch.cuda.automatic_mixed_precision import GradScaler, autocast
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import (
     CosineAnnealingLR,
@@ -105,6 +106,9 @@ def wrapper(self: Trainer[BaseConfig, Any], config: ModelConfigT) -> fl.Module:
             if config.requires_grad is not None:
                 model.requires_grad_(requires_grad=config.requires_grad)
             learnable_parameters = [param for param in model.parameters() if param.requires_grad]
+            if self.config.training.automatic_mixed_precision:
+                for learnable_parameter in learnable_parameters:
+                    learnable_parameter.to(dtype=float32)
             self.models[name] = ModelItem(
                 name=name, config=config, model=model, learnable_parameters=learnable_parameters
             )
@@ -183,6 +187,12 @@ def dtype(self) -> DType:
         logger.info(f"Using dtype: {dtype}")
         return dtype
 
+    @cached_property
+    def scaler(self) -> GradScaler | None:
+        if self.dtype != float16 or not self.config.training.automatic_mixed_precision:
+            return None
+        return GradScaler()
+
     @property
     def learnable_parameters(self) -> list[nn.Parameter]:
         """Returns a list of learnable parameters in all models"""
@@ -341,15 +351,28 @@ 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
+
+    def optimizer_step(self) -> None:
+        if self.scaler is None:
+            self.optimizer.step()
+            return
+        self.scaler.step(self.optimizer)  # type: ignore
+        self.scaler.update()  # type: ignore
+
     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:
@@ -362,7 +385,8 @@ def backward(self) -> None:
     def step(self, batch: Batch) -> None:
         """Perform a single training step."""
         self._call_callbacks(event_name="on_compute_loss_begin")
-        loss = self.compute_loss(batch=batch)
+        with autocast(dtype=self.dtype, enabled=self.config.training.automatic_mixed_precision):
+            loss = self.compute_loss(batch=batch)
         self.loss = loss
         self._call_callbacks(event_name="on_compute_loss_end")
         self.backward()
@@ -403,7 +427,8 @@ def evaluate(self) -> None:
         """Evaluate the model."""
         self.set_models_to_mode(mode="eval")
         self._call_callbacks(event_name="on_evaluate_begin")
-        self.compute_evaluation()
+        with autocast(dtype=self.dtype, enabled=self.config.training.automatic_mixed_precision):
+            self.compute_evaluation()
         self._call_callbacks(event_name="on_evaluate_end")
         self.set_models_to_mode(mode="train")