fix: topkgating major bug

[excepshenal]

topk_masked_gates was previously being used across the tokens dimension to determine which tokens had the highest importance for each expert. However, it was using logits rather than probabilities to determine this.

This was causing print statements like

            if dist.get_rank() == 0:
                print(f"Mask mean: {mask.float().mean()}")
                print(f"Capacity mask mean: {capacity_mask.mean()}")
            mask = torch.logical_and(mask, capacity_mask)
            if dist.get_rank() == 0:
                print(f"Mask (after AND) mean: {mask.float().mean()}")

to often yield values like

Mask mean: 0.0625
Capacity mask mean: 0.0625
Mask (after AND) mean: 0.005908316932618618

and in turn the average number of routed experts per token was as low as 0.001.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 513c3bc413

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[delock]

Hi @excepshenal ，can you fix DCO error by sign-off your commits? Thanks!