The discard mechanism of MoE layer seems wrong.

I am using the deepspeed.moe.layer.MoE. The most concerning part is the dropped inputs. Some inputs are dropped so that experts are not fed in too much data. But when generating outputs in this line. I notice the position holding the discarded inputs are all 0, which is due to the combined_weights are 0. Does not it should be the directly forwarded inputs?

To reproduce the results, run the Cifar example and print the intermediate outputs of the MoE layer. When the inputs are dispatched unevenly, the outputs will see all zero vectors.

[hobbitlzy]

Thanks, @awan-10 @ykim362 . I previously followed the implementation from this repo. They also have a nice implementation but pass the dropped tokens. Anyway, I think you are right. What the Switch Transformer and Gshard recover the dropped tokens by residue connection.

I notice you have achieved similar results as Switch Transformer in this article. I have tried to integrate the MoE module of deepspeed with the training scripts of Huggingface, but I cannot get the performance post by the paper. Actually, my performance is the opposite: the more experts, the lower the sample training efficiency. I find there is one example of MoE on cifar in DeepSpeedExamples. Could you also share the scripts to reproduce the results on Switch Transformers?

[awan-10]

Closing it due to being inactive for a long time. @hobbitlzy – please open if you have more updates on this.