[BUG] ZeRO Stage 2 seems to train MoE models incorrectly

Describe the bug

ZeRO Stage 2 seems to fail updating the weight of MoE gates

To Reproduce

Steps to reproduce the behavior:
we need to injected some code first:

for deepspeed/runtime/engine.py:

1. add self.module (the model) as an input for DeepSpeedZeroOptimizer (about line 1500) for running stage 1 and 2

for deepspeed/runtime/zero/stage_1_and_2.py:

1. in function __init__, add the just mentioned model, I named it as self.module; add self.grad_acc_param = list() to collect ids of param caught by grad hook;

2. in function create_reduce_and_remove_grad_hooks, add self.grad_hooked_param = set() before the whole loop to collect hooked params, right after wrapper(param), add
   param.test_id = len(self.grad_hooked_param) self.grad_hooked_param.add(param.test_id)
   to collect their ids;

3. in function reduce_independent_p_g_buckets_and_remove_grads, add self.grad_acc_param.append(param.test_id) to collect the param caught by grad hooks

4. in function step, inject this:

    diff = set(self.grad_hooked_param) - set(self.grad_acc_param)
    print(len(diff))
    
    import torch
    self.module: torch.nn.Module = self.module
    for module_name, module in self.module.named_modules():
        module: torch.nn.Module = module
        for param_name, param in module.named_parameters(recurse=False):
            if param.test_id in diff:
                print(f"module_name: {module_name}, param_name: {param_name}")
                print(f"param.requires_grad: {param.requires_grad}")
   

to check if any require grad param is not caught by grad hook

run the code

run official Megatron-DeepSpeed/examples_deepspeed/MoE/ds_pretrain_gpt_125M_MoE64.sh

Expected behavior

nothing should be printed (or all printed params require no grads)

actual output

module_name: language_model.encoder.layers.1.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True
module_name: language_model.encoder.layers.3.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True
module_name: language_model.encoder.layers.5.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True
module_name: language_model.encoder.layers.7.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True
module_name: language_model.encoder.layers.9.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True
module_name: language_model.encoder.layers.11.mlp.deepspeed_moe.gate.wg, param_name: weight
param.requires_grad: True

ds_report

Setup is included in Megatron-DeepSpeed/examples_deepspeed/MoE/ds_pretrain_gpt_125M_MoE64.sh

System info (please complete the following information):

• OS: Debian GNU/Linux 11 (bullseye)
• GPU count and types: one machines with x4 A40
• Python version 3.8.8

Start in deepspeed/runtime/engine.py around the DeepSpeedZeroOptimizer construction and deepspeed/runtime/zero/stage_1_and_2.py in init, create_reduce_and_remove_grad_hooks, reduce_independent_p_g_buckets_and_remove_grads, and step. Reproduce with Megatron-DeepSpeed's examples_deepspeed/MoE/ds_pretrain_gpt_125M_MoE64.sh and inspect whether the listed MoE gate weights receive gradient hooks and updates. Done means the reported trainable gate parameters are handled correctly without the diagnostic output.