diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 392dcf8..1d384c9 100644 --- a/ScaFFold/utils/trainer.py +++ b/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.")