github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Allow automagic3 to run unfused. Add some clipping.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-08 08:59:06 -06:00
parent c9cdbb5bb7
commit e281df70dd
2 changed files with 89 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/optimizers/test_optimizers.py
View File

@@ -8,7 +8,7 @@ Compares speed (ms/step) and peak VRAM across:
- Adafactor
- Automagic v1
- Automagic v2 (fused-backward)
- Automagic v3 (fused-backward)
- Automagic v3 (fused-backward and traditional/unfused)
- Prodigy
"""
import contextlib
@@ -186,7 +186,8 @@ def main():
("Adafactor", lambda p: Adafactor(p, lr=1e-4, scale_parameter=False, relative_step=False, warmup_init=False)),
("Automagic v1", lambda p: Automagic(p, lr=1e-4)),
("Automagic v2", lambda p: Automagic2(p, lr=1e-4)),
("Automagic v3", lambda p: Automagic3(p, lr=1e-4)),
("Automagic v3 fused", lambda p: Automagic3(p, lr=1e-4, fused=True)),
("Automagic v3 unfused", lambda p: Automagic3(p, lr=1e-4, fused=False)),
("Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)),
]