[HELP] How to safely switch trainable parameters in ZeRO-3 stage? about deepspeed HOT 2 CLOSED

I managed to solve this issue by

• remove the previous gradient synchronization hooks
• reset the bookkeeping of ds optimizer
• setup the fp16 groups and partition it
• register necessary hooks for synchronizing gradients
• create fp32 flat partition, initialize ipg buffer and grad partition buffer
• invalidate the trace cache, since the module processing order has been changed

after switching the trainable parameters each time. This solution is more time-efficient compared to rerun deepspeed.initialize each time. Interested readers may refer to this repo, where i implement the main logics in _switch_trainable_params_zero3.

from deepspeed.

Thanks for adding your solution @Ledzy

from deepspeed.