How the experts' gradients are handled under data parallelism?

[yzs981130]

When count_per_node is set to negative, one expert should be paralleled on multiple GPUs like ZERO, with each GPU holding a slice of the expert's parameters. There are also all_gathers performed within the ffn_zero_group in the forward of the expert.

My question is how these gradients and parameter update of the expert is handled in Tutel under DP. Examples seem to indicate that there are no extra efforts for users to manually handle it. However, I cannot find the corresponding implementation in either moe_layer or TutelDistributedOptimizer.

Any help will be appreciated!

[ghostplant]

You can reference the implementation here: https://github.com/microsoft/tutel/blob/main/tutel/experts/ffn.py