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.
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 alsoall_gather
s performed within theffn_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
orTutelDistributedOptimizer
.Any help will be appreciated!