Enable compatibility of sparse MOE training under DeepSpeed ZeRO3

[Jacky-hate]

According to deepspeedai/DeepSpeed#4966, ZeRO3 in DeepSpeed does not work with MoE models because the order of executing modules can change at every forward/backward pass and a new API is implemented to stop breaking down a module for parameter fetching. Similar case occurs when finetuning Qwen1.5-MoE-A2.7B using ZeRO3 optimization(#275).

This PR use the api above to make sparse MoE layer compatible with Zero3.

[jklj077]

Hi, thanks for your contribution! It is a very nice addition.

Unfortunately, we are in the process of deprecating the examples/sft/finetune.py script and the related due to reported compatibility issues. As our code is adapted from https://github.com/lm-sys/FastChat/tree/main/fastchat/train, I suppose the same issue also exists there and they may still be open to PR on the training code.

thanks again for the support.