Why is the ZeRO3 optimizer only capable of certain types? Any special reason?

[yhna940]

Why is the ZeRO3 optimizer only capable of certain types? Any special reason?

ColossalAI/colossalai/zero/gemini/gemini_optimizer.py

Line 72 in ffcdbf0

assert type(optim) in _AVAIL_OPTIM_LIST, "You should use an optimizer in the available list:\n" \

What can I do if I want to use this optimizer more generically? (I also want to use optimizers like RMS or SGD.)

[1SAA]

Hi, @yhna940
The reason is that the gradient generated is FP16, and the master weight is FP32. We just need a optimizer that can accept FP16 gradients with FP32 parameters.

[yhna940]

@1SAA
Thank you for answer.
Can't we just cast the gradient to the master parameter's type, like low level zero?
I would be grateful if you could provide feedback on my proposed changes.
#3699