[REQUEST] Non-element-wise Optimizer Compatibility

[Triang-jyed-driung]

I am encountering issues when using non-element-wise optimizers such as Adam-mini with DeepSpeed.
According to the documentation, it reads:

The FP16 Optimizer is designed to maximize the achievable memory bandwidth by merging all the parameters of the model into a single large buffer, and applying the weight updates in a single kernel, allowing it to achieve high memory bandwidth.

As for Adam-Mini, the merging of parameters into a single buffer degenerates it to SGD with only one global second moment. This will corrupt the training, causing spikes and NaNs.

Recently have seen a trend developing non-element-wise optimizers. Besides Adam-Mini, there are several different ones using matrix SVD, like Fira (https://arxiv.org/abs/2410.01623) and Muon (https://github.com/KellerJordan/modded-nanogpt/blob/master/train_gpt2.py). Blindly merging parameters will corrupt these optimizers as well. This limitation is impacting my ability to fully leverage DeepSpeed for my training needs.

Of course, this issue can be avoided by using naive DDP in Lightning. But the speed also slows down unbearably by 40%. Is there a recommended approach or workaround for integrating non-element-wise optimizers with DeepSpeed? Or are there plans to support such optimizers in future releases?

Thank you for your assistance.

[tjruwase]

@Triang-jyed-driung, can you clarify if you are interested in integration with ZeRO specifically or DeepSpeed in general?
Also, can you take a look at our unfused Adam optimizer: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/fp16/unfused_optimizer.py? How different is this from your proposed Adam-Mini? Thanks!

[Triang-jyed-driung]

My optimizer is a simplified version of https://github.com/zyushun/Adam-mini that shares one second moment for each matrix/vector (no splitting between key/value/heads), which is simpler and (theoretically) faster. But still, the model is split into approximately 450 parameter groups, so fusing becomes impossible.
However, the performance dropped by nearly 40% compared to Adam, from 60 kt/s to just over 30 kt/s. It's likely because that Python has to loop over these parameters one by one in a single optimization step.
I use default strategy of "deepspeed_stage_[1,2,3]" of Pytorch Lightning, so I guess it's related to Deepspeed Zero.