Use fused Adam - #104
Open
ndryden wants to merge 1 commit into
Open
Conversation
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.
michaelmckinsey1
approved these changes
Aug 13, 2026
Comment on lines
+256
to
+277
| # 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. |
Collaborator
There was a problem hiding this comment.
comment necessary? maybe a sentence or two is warranted.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This gives a small performance improvement, but it's very low-hanging fruit.
Code by Claude.