From d15c0d591818dc2a62ba6e712d87cab72f76d03f Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 22:49:08 -0700 Subject: [PATCH 1/2] 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 392dcf8..f99ba39 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.") From 2ad606aa6bcfe5ed015071f8571aae6fc877178a Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Tue, 8 Sep 2026 18:44:50 -0700 Subject: [PATCH 2/2] Tighten the fused-Adam comment and drop its measured figures Reviewers asked for more concise comments without performance numbers that go stale. The comment now keeps the reasons -- one kernel instead of several, the optimizer is the line item sharding does not shrink, CUDA only, fused and foreach are not bitwise comparable with each other but each is reproducible with itself including across checkpoint/resume -- and drops the per-scale timings, the step deltas, and the run counts. --- ScaFFold/utils/trainer.py | 29 +++++++++-------------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index f99ba39..1d384c9 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -253,28 +253,17 @@ def setup_training_components(self): """Set up the optimizer, scheduler, gradient scaler, and loss function.""" # Set up optimizer if self.config.optimizer == "ADAM": - # 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. + # 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. 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. + # 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(