setup fabric

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -9,6 +9,7 @@
 
 from ema_pytorch import EMA
 
+from lightning import Fabric
 
 # helpers
 
@@ -49,9 +50,21 @@ def __init__(
         ),
         clip_grad_norm = 10.,
         default_lambda_lr = default_lambda_lr_fn,
+        fabric: Fabric | None = None,
+        accelerator = 'auto',
+        fabric_kwargs: dict = dict(),
         ema_kwargs: dict = dict()
     ):
         super().__init__()
+
+        if not exists(fabric):
+            fabric = Fabric(accelerator = accelerator, **fabric_kwargs)
+
+        self.fabric = fabric
+        fabric.launch()
+
+        # model
+
         self.model = model
 
         # exponential moving average
@@ -74,6 +87,10 @@ def __init__(
 
         self.optimizer = optimizer
 
+        # setup fabric
+
+        self.model, self.optimizer = fabric.setup(self.model, self.optimizer)
+
         # scheduler
 
         if not exists(scheduler):

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
     "ema-pytorch>=0.4.8",
     "environs",
     "jaxtyping>=0.2.28",
-    "pytorch-lightning>=2.2.5",
+    "lightning>=2.2.5",
     "taylor-series-linear-attention>=0.1.9",
     "torch>=2.1",
     "tqdm",