Compiled basic optimizers

[dxqb]

Draft implementation

• add UI
• AdamW not planned
• coordinate with Compiled Optimizers: Accelerate all Advanced Optimizers with pre-compilation #1083 and https://github.com/Koratahiu/Advanced_Optimizers/tree/compile_2

[maedtb]

Should this and the below step_parameter function get decorators like @functools.wraps(step_adafactor) to make these wrappers (slightly more) seamless?

[dxqb]

can you elaborate? what's the problem with the current code?

[dxqb]

AdamW_adv can do everything what AdamW can but also compile, so there is not much point of implementing it again for patched AdamW.

However, Adafactor is the preset for full finetuning and compile has a large effect in full finetuning. adv_optm currently doesn't have Adafactor
CC @Koratahiu

[Koratahiu]

However, Adafactor is the preset for full finetuning and compile has a large effect in full finetuning. adv_optm currently doesn't have Adafactor CC @Koratahiu

Any adv Adam variant with beta1 = 0 and the factored option is essentially Improved Adafactor (factored second moment, but with uncompressed raw gradient contribution and better factorization ).
I am also implementing a factored 2nd moment option in #1344, which behaves just like Adafactor (dense first moment and factored second moment).

I guess you can do factored AdamW_adv with beta1=0 for full finetuning presets?