[REQUEST] Non-element-wise Optimizer Compatibility

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.

Start with the FP16 Optimizer documentation and the parameter-merging behavior described in the issue. Compare that behavior with Adam-Mini, Fira, and Muon, then investigate how DeepSpeed integrates optimizers during distributed training. Done means establishing a supported approach for non-element-wise optimizers without corrupted updates, spikes, or NaNs.