[BUG] DeepSpeed-MoE TopK gate : Is it supposed to not be trained?

While running the MoE cifar10 example, I added a small code in the hook function for CPU-Offload version to see which parameters are backpropagated during backward function. Find this code here, when registering hooks to the gradient accumulation function of each parameter.

def reduce_partition_and_remove_grads(*notneeded):
    print("Running hook for {}".format(self.get_param_id(param))) # Added code
    self.reduce_ready_partitions_and_remove_grads(param, i)

My results shoes that parameter with param_id 8 is not back-propagated (missing hook for 8)

Running hook for 10
Running hook for 9
Running hook for 14
Running hook for 13
Running hook for 12
Running hook for 11
Running hook for 7
Running hook for 6
Running hook for 5
Running hook for 4
Running hook for 2
Running hook for 3
Running hook for 0
Running hook for 1

Then I used this code to show the param_id for each parameter. This was run in the main training code.

for name,param in model_engine.module.named_parameters():
    print(name,param.numel(),model_engine.optimizer.get_param_id(param))

My results showed that the parameter with param_id 8, i.e. the TopK gating layer was not back-propagated during backward() call :

conv1.weight 450 0
conv1.bias 6 1
conv2.weight 2400 2
conv2.bias 16 3
fc1.weight 48000 4
fc1.bias 120 5
fc2.weight 10080 6
fc2.bias 84 7
moe_layer_list.0.deepspeed_moe.gate.wg.weight 168 8  # Backward pass was not called
moe_layer_list.0.deepspeed_moe.experts.deepspeed_experts.0.weight 7056 11
moe_layer_list.0.deepspeed_moe.experts.deepspeed_experts.0.bias 84 12
moe_layer_list.0.deepspeed_moe.experts.deepspeed_experts.1.weight 7056 13
moe_layer_list.0.deepspeed_moe.experts.deepspeed_experts.1.bias 84 14
fc4.weight 840 9
fc4.bias 10 10

I was wondering if this was an expected behavior.
I checked and confirmed requires_grad=True for this particular parameter.