Repository: huggingface/transformers
Pull request: Fix broken Llama4 accuracy in MoE part

Llama4 accuracy is broken by a bug in
https://github.com/huggingface/transformers/pull/39501 . It forgot to transpose the router_scores before applying it to routed_in, causing Llama4 to generate garbage output.

This PR fixes that issue by adding back the transpose() and adding some comments explaining why the transpose() is needed.

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

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.

@ArthurZucker please review since https://github.com/huggingface/transformers/pull/39501 was made by you. Thanks!

## Accuracy tests

Test script:

```
from transformers import pipeline
import torch

model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

pipe = pipeline(
    "text-generation",
    model=model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

output = pipe("Roses are red,", max_new_tokens=200)

print(output)
```

Before the fix on H200:

```
[{'generated_text': 'Roses are red, 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 
8 8 8 8 8 8 8 8 8 8 8 8 8 8 8'}]
```

After the fix on H200:

```
[{'generated_text': "Roses are red, violets are blue, and here are some Valentine's Day-themed books to read with your boo! \n Celebrate Black History Month by reading books by Black authors and about Black culture!  These books are for kids and teens! \n These picture books celebrate Black History Month and are perfect for reading with your kids!  These books highlight Black history, cu... \n These books are all winners or nominees for major literary awards, including the Pulitzer Prize, the National Book Award, and t... \n 100 Book Challenge 2024: Read a Book Set in Another Country \n Check out these books that are set in different countries around the world!  Read a book set in another country and see what you... \n Explore the diverse world through these books that are set in different countries and cultures!  Read a book set in another coun... \n 100 Book Challenge 2024: Read a Book Written by an Author of Color \n Check out these books written by authors of color!"}]
```

Changed files:
filename: src/transformers/models/llama4/modeling_llama4.py
patch: @@ -158,7 +158,7 @@ def forward(self, hidden_states):
         hidden_states = hidden_states.reshape(-1, self.hidden_dim)
         router_scores, router_logits = self.router(hidden_states)
         routed_in = hidden_states.repeat(router_scores.shape[1], 1)
-        routed_in = routed_in * router_scores.reshape(-1, 1)
+        routed_in = routed_in * router_scores.transpose(0, 1).reshape(-1, 1)
         routed_out = self.experts(routed_in)
         out = self.shared_expert(hidden_states)
         out.add_(routed_out.reshape(router_scores.shape[1], -1, routed_out.shape[-1]).sum(dim=0))