Use fused Adam

ScaFFold/utils/trainer.py

@@ -253,9 +253,23 @@ def setup_training_components(self):
         """Set up the optimizer, scheduler, gradient scaler, and loss function."""
         # Set up optimizer
         if self.config.optimizer == "ADAM":
-            self.log.info("Using ADAM optimizer.")
+            # Fused Adam does the whole update in one kernel instead of the
+            # foreach path's several.  Parameters are replicated across ranks,
+            # so the optimizer is the one line item spatial sharding does not
+            # shrink; this saving lands whole on every rank.
+            #
+            # CUDA only: the fused kernels are device-specific, and the CPU
+            # trainers the tests build have nothing to gain.  Fused and foreach
+            # accumulate in different orders, so runs are not bitwise
+            # comparable across the two; each is reproducible with itself,
+            # including through checkpoint/resume, where the fused path keeps
+            # its ``step`` counter on the device.
+            fused = self.device.type == "cuda"
+            self.log.info(f"Using ADAM optimizer{' (fused)' if fused else ''}.")
             self.optimizer = optim.Adam(
-                self.model.parameters(), lr=self.config.starting_learning_rate
+                self.model.parameters(),
+                lr=self.config.starting_learning_rate,
+                fused=fused,
             )
         elif self.config.optimizer == "SGD":
             self.log.info("Using SGD optimizer.")