Adam-Mini

[lannashelton]

A new optimizer called Adam-Mini was released a few days ago. Are we going to see this added to Kohya-SS? From what it claims, it appears to reduce memory usage about 50% and reduces training time about 33%

https://arxiv.org/abs/2406.16793

[5KilosOfCheese]

From the paper:

"Higher throughput. Adam-mini can reach a higher throughput than AdamW, especially under limited hardware resources. There are two reasons. First, thanks to the memory cut-down, Adam-mini largely reduces the communication among CPUs and GPUs, which is known to be a major overhead [ 42 ]. Second, Adam-mini does not introduce extra computation in its update rules. The average operation incurs negligible cost. Additionally, it significantly reduces the number of tensor-square-root and tensor-division operations in AdamW. We further highlight that Adam-mini does not conduct matrix factorization, which is expensive but often occurs in low-rank-approximation based optimizer (e.g. [59]). Adam-mini achieves faster per-step updates than these methods." (From the paper).

And the performance is basically the same.

Thats actually... Like a really big deal. Currently all AI model development is basically just hogging all the resources you can, but this working allows smaller players (like community driven projects) to have a more realistic chance at participating.

This actually opens interesting options. Like lets play around with this thought. FP16 cuts thing even more, then possibly FP8 yet again half. And since we can gain some performance at each step. They address this at the end of the paper:
"There are several other tricks that ease GPU memory burden but are orthogonal to optimizer design. These tricks include gradient checkpointing [5], model offload and sharding [ 41, 42 ], quantization [11, 25], and fused update [31, 32]. Adam-mini can be implemented upon these tricks."

This would actually make things even more accesable. Very exciting.