Update sharded_moe.py to support top2 gate with Tutel

[xenshinu]

Tutel is forced to be unused on k > 1 since #2053
Given the fact that multiple experts per token is very common, and the gather and scatter operation without Tutel is so inefficient, I added support of tutel to top2 gate and tested on pipeline engine. This can be done for any k actually, I'll push that later when I have time to test,

[xenshinu]

@xenshinu please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree company="University of Michigan"

[xenshinu]

Not sure if this check is needed

    if use_tutel:
        # Tutel doesn't support index values masked with zero
        # so we need to replace masked indices with -1
        indices1_mask = mask1.sum(dim=1) * num_experts - 1
        indices1_s = torch.min(indices1_s, indices1_mask)
        indices2_mask = mask2.sum(dim=1) * num_experts - 1
        indices2_s = torch.min(indices2_s, indices2_mask)

I see that in top1gate,

DeepSpeed/deepspeed/moe/sharded_moe.py

Line 252 in 66d3d3e

if use_tutel:

but when I refer to examples from Tutel
https://github.com/microsoft/Tutel/blob/ab7937bb929bc78111d74261b490da25657a7e5c/tutel/impls/fast_dispatch.py#L143

I didn't see any specify check for non-zero mask.

[loadams]

@xenshinu - thanks for this PR, could you run the pre-commit formatter on your branch to resolve the "Formatting" error? I believe it just wants the use_tutel on a new line here.

[xenshinu]

Thanks for pointing that out. I've updated the file and it looks like the pre-commit check has passed.
Can someone take a look to this question?
#6948 (comment)