I have a node with 8 GPUs. The model has 8 experts and I use TP=1, EP=8, with sequence parallel on. Then I expect each GPU has one expert. I use the MoEAllGatherTokenDispatcher. The size of the hidden_states is [S/TP, B, H] for token_permutation, which is actually [S, B, H] because TP=1. Then why we still need tensor_parallel.gather_from_sequence_parallel_region_to_moe to gather and form a global_hidden_states, whose size will become [S * EP, B, H]? In my view, each rank has a copy of the [S, B, H] hidden_states, why there is still a need for all gather?
There are only B * S tokens to compute, but now each rank has B * S * EP tokens, because we use get_tensor_and_expert_parallel_group() and its size is TP * EP=8.
I have a node with 8 GPUs. The model has 8 experts and I use TP=1, EP=8, with sequence parallel on. Then I expect each GPU has one expert. I use the
MoEAllGatherTokenDispatcher. The size of the hidden_states is [S/TP, B, H] fortoken_permutation, which is actually [S, B, H] because TP=1. Then why we still needtensor_parallel.gather_from_sequence_parallel_region_to_moeto gather and form aglobal_hidden_states, whose size will become [S * EP, B, H]? In my view, each rank has a copy of the [S, B, H] hidden_states, why there is still a need for all gather?There are only B * S tokens to compute, but now each rank has B * S * EP tokens, because we use get_tensor_and_expert_parallel_group() and its size is TP * EP=8.