Fix load balancing loss func for mixtral

[liangxuZhang]

What does this PR do?

Fixes #28255

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker and @younesbelkada
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[kalomaze]

How does this differ from #28115 ?

[ArthurZucker]

Thanks a lot for deep diving. As discussed here this is welcome, we indeed had a bug in the implementation

Let's try to help with shapes and use something explicit like this:

    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1) # [batch_size X sequence_length, top_k]

    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) # [batch_size X sequence_length, top_k, num_experts]

    tokens_per_expert = torch.mean(expert_mask.float(), dim=0) # [top_k, num_experts]

    # Compute the average probability of routing to these experts
    router_prob_per_expert = torch.mean(routing_weights, dim=0) # [num_experts]

    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0)) / top_k
    return overall_loss * num_experts

[ArthurZucker]

this is not needed if we just remove selected_experts = selected_experts.reshape(-1)

[ArthurZucker]

Suggested change

expert_mask = torch.max(expert_mask, dim=-2).values

should be removed

[codybum]

What is the impact of this issue on Mixtral training? Will this fix conceivability improve the quality of training? Is it likely that previous Mixtral trainings are not as good as they could be?

It seems like an important issue for those working with Mixtral that has been waiting on merge approval for a while.

[theblackcat102]

I personally finds the loss to be much lower with the new implementation. But I wasn't sure if it has to do with the (num_experts**2) instead of just N. I'm pretty sure this is an error on original mixtral side. So far still waiting for the training result on new implemented balance loss to finish. Deepspeed also has an implementation of top-2 which we might be able to reference.

[ArthurZucker]

#28255 has information that could help, I am down to merge this for the release planned this week, just the comments that need to be adressed cc @liangxuZhang do you need help to finish this?

[liangxuZhang]

#28255 has information that could help, I am down to merge this for the release planned this week, just the comments that need to be adressed cc @liangxuZhang do you need help to finish this?

@ArthurZucker LGTM. The new implementation is correct and concise, and I've made a new commit. In #28255, maybe we can have a deep discuss whether to concatenate gate logits of all layers.

[ArthurZucker]

Alright! Pretty sure the math shows it's equivalent to compute on individual layers then sum vs concate and compute, but let's merge this for now !

[ArthurZucker]

Thanks! Failing test seems unrelated let's just rebase on main

[liangxuZhang]

Thanks! Failing test seems unrelated let's just rebase on main

@ArthurZucker I've just rebase on the main branch, but I'm not sure if I'm doing it right. Please tell me what else I need to do

[bratao]

@liangxuZhang @ArthurZucker opinions about #28403 ? It looks complementary to this PR

[ArthurZucker]

Something like git pull upstream main if the remote is upstream, the exotic CI was fixed on main! I'll merge without it

[ArthurZucker]

Thanks a lot @liangxuZhang for this fix! 🤗

[dancingpipi]

great！

[cryoco]

Impressive work!