Hey! The generate function is not supposed to work for training. That is why we don't test past key values and output router logits. Though it's actually not that incompatible (you could want to look at the distribution of the router logits during generation).
Do you want to open a PR for a fix?

Hi, could I take this up? @ArthurZucker

Sur feel free to open a PR!

I am having trouble figuring out how to go about the [batch_size X 1, num_experts] case. I am thinking of creating a condition for this in the load_balancing_loss_func, such as, is_single_token = gate_logits[0].shape[0] == 1. After that, I will modify the code in the following way, so that the shapes during the multiplication are valid:

        if not is_single_token:
            expert_attention_mask = (
                attention_mask[None, :, :, None, None]
                .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
                .reshape(-1, top_k, num_experts)
                .to(compute_device)
            )
        else:
            expert_attention_mask = (
                attention_mask[None, :, -1:, None, None]
                .expand((len(gate_logits), 1,1, top_k, num_experts))
                .reshape(-1, top_k, num_experts)
                .to(compute_device)
            )

As you can see here, I think one approach is to retain the attention score of that particular token itself (the one that we are adding), and get rid of the other scores in the sequence. This is because, if we do retain the other ones, then expert_attention_mask's shape will not match that of expert_mask. As a consequence, for this case, we will not need to calculate the num_hidden_layers (as it is a tiny value as mentioned by @sorgfresser) , because we will just get the batch size directly from gate_logits. Similarly, we will need a modification for router_per_expert_attention_mask:

if not is_single_token:
            router_per_expert_attention_mask = (
                attention_mask[None, :, :, None]
                .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
                .reshape(-1, num_experts)
                .to(compute_device)
            )
        else:
            router_per_expert_attention_mask = (
                attention_mask[None, :, -1:, None]
                .expand((len(gate_logits), 1, 1, num_experts))
                .reshape(-1, num_experts)
                .to(compute_device)
            )

This approach is what I have come up with, I hope you all can verify the correctness of the same, and point out if I have made any mistakes. The PR is here.