Ethan @ethansmith2000.com · Nov 23 Here, have PSGD-Kron and SOAP with FSDP2 support. Please go wild with it, let's see something finally replace ADAM. github.com/ethansmith20...
13 likes 3 replies
?
Replies
What is the memory profile like on these optims?
This is only kron whitening btw. We should port the Hessian version too.
naive implementation is okay, much needs to be done with full tensor but a few pointwise ops we can benefit from remaining in shards. future improvements likely lie in pre-fetching and grabbing chunks of params, rather than one at a time. this is fragile though and DTensor still needs more support.