From 9478fba499a02ce59d41b95d499f49edd58f27a1 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 22:49:08 -0700 Subject: [PATCH] Run Adam fused on CUDA Seven ``_foreach_`` launches collapse into one ``_fused_adam_``, taking the optimizer line item from 3.199 to 1.617 ms at scale 7 and from 11.999 to 5.884 at scale 8 (kernel-only device time). End to end, paired arms alternating within each rep, 6 reps, single GPU: scale 7 (128^3) 66.32 -> 65.09 ms/step -1.23 +/- 0.80 0.9815x scale 8 (256^3) 449.74 -> 443.57 ms/step -6.17 +/- 0.53 0.9863x 12/12 pairs the same sign, peak memory unchanged to the digit. At scale 8 the accounting closes: total device time -6.011 ms against -6.162 in the optimizer, everything else netting +0.151. The reason it is worth a commit at that size is where the time sits. The optimizer is the one line item in the step breakdown that does not shrink with spatial sharding -- 12.0 ms at scale 8 on one, two and four GPUs alike, so it grows from 2.7% of a step at 1 GPU to 7.1% at 4 -- which means this saving lands whole on every rank instead of being divided among them. CUDA only. The fused kernels are device-specific and the CPU trainers the tests build have nothing to gain, so the flag is derived from the trainer's own device rather than assumed. Not numerically free: fused Adam accumulates differently from foreach and moves the loss by up to 6.7e-6 relative over a run, so runs across this commit are not bitwise comparable. Each arm is still reproducible with itself, which is the property this branch has been protecting, and that was measured rather than assumed -- 7 independent scale-7 runs (4 default, 3 more_determinism=1) bitwise identical in parameters, per-batch loss and dice, forward activations and every train_stats.csv column bar wall-clock, with default still equal to more_determinism=1 bit for bit. A foreach control on this same tree is equally reproducible and differs from fused only downstream of the first optimizer step, which places the numerics change where it belongs. Checkpoint save and resume are bitwise transparent, verified at production volume. That is not free here: the fused path keeps its step counter on the device where foreach keeps it on the host, and checkpointing.py moves optimizer state to CPU to save it. Suite unchanged at 752 passed / 8 skipped / 1 xfailed. Untested: the GradScaler interaction. It is disabled under bf16, so _fused_adam_'s found_inf path is not exercised by any run behind these numbers. --- ScaFFold/utils/trainer.py | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 201e09f..8fa4075 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -253,9 +253,34 @@ 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.") + # The fused path does the whole parameter update in one kernel + # rather than the foreach path's several: seven ``_foreach_`` + # launches collapse into one ``_fused_adam_``, taking the optimizer + # line item from 3.199 to 1.617 ms at scale 7 and 11.999 to 5.884 + # at scale 8 (kernel-only device time). End to end, paired arms + # alternating within each rep, 6 reps: **-1.23 +/- 0.80 ms/step at + # scale 7 and -6.17 +/- 0.53 at scale 8**, 12/12 pairs same sign. + # + # Worth having because the optimizer is the one line item that does + # *not* shrink with spatial sharding -- 12.0 ms at scale 8 on 1, 2 + # and 4 GPUs alike, so it grows from 2.7% of a step to 7.1% as + # ranks are added -- which means this saving lands whole on every + # rank instead of being divided among them. + # + # CUDA only: the fused kernels are device-specific, and the CPU + # trainers the tests build have nothing to gain. Not free + # numerically -- fused Adam accumulates differently from foreach, + # which moves the loss by up to 6.7e-6 relative over a run. Each + # arm is still reproducible with itself, measured: 7 independent + # scale-7 runs bitwise identical, and checkpoint/resume bitwise + # transparent, which is not free here because the fused path keeps + # its ``step`` counter on the device rather than on the host. + 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.")