search

microsoft / tutel

Tutel MoE: An Optimized Mixture-of-Experts Implementation
MIT License
694 stars 84 forks source link

How the experts' gradients are handled under data parallelism? #192

Open yzs981130 opened 1 year ago

yzs981130 commented 1 year ago

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 commented 1 year ago

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