Use fused Adam - #104
Merged
Merged
Use fused Adam#104
Conversation
michaelmckinsey1
approved these changes
Aug 13, 2026
PatrickRMiles
approved these changes
Aug 27, 2026
PatrickRMiles
left a comment
PatrickRMiles
left a comment
Collaborator
There was a problem hiding this comment.
Agreed that details of the experimentation Fable performed to evaluate this change should not be included in a comment in the code. If we even want those details hanging around somewhere (arguable), a commit message is probably best.
ndryden
force-pushed
the
fused-adam
branch
2 times, most recently
from
September 18, 2026 21:40
aaa5295 to
cda0676
Compare
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.
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.
Reviewers found the comment still too long. It now states the saving, the CUDA-only condition, and the reproducibility contract in one paragraph. Comment only; the code is unchanged.
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
3 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.