Make float16 actually work using gradient and optimizer state scaling - #15

Open

e-c-d wants to merge 4 commits into

warner-benjamin:mainfrom

e-c-d:feature/gradient-scaling

Commits on Mar 23, 2026

Implement gradient and optimizer state scaling for Adam and AdamW, excluding coupled weight decay. Test cases now pass for float16.
e-c-d
committed
9138f0c

Commits on Mar 24, 2026

Implement gradient and optimizer state scaling for StableAdamW for the triton backend, allowing the tests to pass using float16.
e-c-d
committed
2d9e9f9

Commits on Mar 25, 2026