Fix broken Llama4 accuracy in MoE part (#40609)
* 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.
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
* remove comment
---------
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>
