Enabled Qwen2-MoE Tensor Parallelism (TP) inference

[gyou2021]

Modified _replace_module in auto_tp.py :
The modification keeps the layers 'shared_expert_gate' and 'gate' in qwen2-moe the original type torch.nn.Linear and not changes them into LinearLayer. In this way, their weights will not be split into multiple HPU/GPU cards. Then the qwen2-moe can run on multiple HPU/GPU cards.
Since the weights of 'gate' are not split into multiple HPU/GPU cards, all gather operations are not needed, which may improve performance.

[delock]

Hi @Yejing-Lai , do you want to provide some comments on this PR for Qwen2-MoE AutoTP support?

[Yejing-Lai]

Could you try to modify this line if it can meet your needs? https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/auto_tp.py#L336

[gyou2021]

Could you try to modify this line if it can meet your needs? https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/auto_tp.py#L336

Yes. It can provide the same function and result if probably coded.

[gyou2021]

Could you try to modify this line if it can meet your needs? https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/module_inject/auto_tp.py#L336

Thank you for your comments.
I just moved the linear filter of qwen2-moe from _replace_module() to _replace() for uniform code management. Both have the same function and the same result. The qwen2-moe branch has been updated.

[delock]

Hi @gyou2021 , can you also add qwen2-moe to this list? Some user will check this page for AutoTP supported model.
https://github.com/microsoft/DeepSpeed/blob/master/docs/_tutorials/automatic-tensor-parallelism.md#supported-models

[gyou2021]

Hi @gyou2021 , can you also add qwen2-moe to this list? Some user will check this page for AutoTP supported model. https://github.com/microsoft/DeepSpeed/blob/master/docs/_tutorials/automatic-tensor-parallelism.md#supported-models

Added. Thank you for your comment.

[delock]

Hi @tjruwase This PR adds AutoTP support for Qwen2-MoE. @Yejing-Lai and me had reviewed this change. Thanks!