add distributed evaluation for fabric training, make sure it can be t… (#291)

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -171,6 +171,7 @@ def __init__(
         checkpoint_folder: str = './checkpoints',
         overwrite_checkpoints: bool = False,
         fabric_kwargs: dict = dict(),
+        distributed_eval: bool = True,
         fp16: bool = False,
         use_ema: bool = True,
         ema_kwargs: dict = dict(
@@ -201,10 +202,16 @@ def __init__(
         self.fabric = fabric
         fabric.launch()
 
+        # whether evaluating only on root node or not
+        # to save on each machine keeping track of EMA
+
+        self.distributed_eval = distributed_eval
+        self.will_eval_or_test = self.is_main or distributed_eval
+
         # exponential moving average
 
         self.ema_model = None
-        self.has_ema = self.is_main and use_ema
+        self.has_ema = self.will_eval_or_test and use_ema
 
         if self.has_ema:
             self.ema_model = EMA(
@@ -282,16 +289,18 @@ def __init__(
         self.valid_every = valid_every
 
         self.needs_valid = exists(valid_dataset)
+        self.valid_dataloader = None
 
-        if self.needs_valid and self.is_main:
+        if self.needs_valid and self.will_eval_or_test:
             self.valid_dataset_size = len(valid_dataset)
             self.valid_dataloader = DataLoader_(valid_dataset, batch_size = batch_size)
 
         # testing dataloader on EMA model
 
         self.needs_test = exists(test_dataset)
+        self.test_dataloader = None
 
-        if self.needs_test and self.is_main:
+        if self.needs_test and self.will_eval_or_test:
             self.test_dataset_size = len(test_dataset)
             self.test_dataloader = DataLoader_(test_dataset, batch_size = batch_size)
 
@@ -306,6 +315,12 @@ def __init__(
 
         fabric.setup_dataloaders(self.dataloader)
 
+        if exists(self.valid_dataloader) and self.distributed_eval:
+            fabric.setup_dataloaders(self.valid_dataloader)
+
+        if exists(self.test_dataloader) and self.distributed_eval:
+            fabric.setup_dataloaders(self.test_dataloader)
+
         # scheduler
 
         if not exists(scheduler):
@@ -555,7 +570,7 @@ def __call__(
             # maybe validate, for now, only on main with EMA model
 
             if (
-                self.is_main and
+                self.will_eval_or_test and
                 self.needs_valid and
                 divisible_by(self.steps, self.valid_every)
             ):
@@ -585,6 +600,11 @@ def __call__(
 
                 valid_loss_breakdown = {f'valid_{k}':v for k, v in valid_loss_breakdown.items()}
 
+                # reduce valid loss breakdown
+
+                if self.distributed_eval:
+                    valid_loss_breakdown = self.fabric.all_reduce(valid_loss_breakdown, reduce_op = 'sum')
+
                 # log
 
                 self.log(**valid_loss_breakdown)
@@ -598,7 +618,7 @@ def __call__(
 
         # maybe test
 
-        if self.is_main and self.needs_test:
+        if self.will_eval_or_test and self.needs_test:
             eval_model = default(self.ema_model, self.model)
 
             with torch.no_grad(), to_device_and_back(eval_model, self.device):
@@ -625,6 +645,11 @@ def __call__(
 
             test_loss_breakdown = {f'test_{k}':v for k, v in test_loss_breakdown.items()}
 
+            # reduce
+
+            if self.distributed_eval:
+                test_loss_breakdown = self.fabric.all_reduce(test_loss_breakdown, reduce_op = 'sum')
+
             # log
 
             self.log(**test_loss_breakdown)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.5.41"
+version = "0.5.42"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },