[BugFix] Fix Llama4 Calibration

[dsikka]

SUMMARY:

• Applying the router_scores to the hidden states before passing the hidden states to the experts is resulting in NaNs during calibration.
• I have gone through the forward pass line-by-line, verified that the dimenions all match / make sense and ensured we are not doing anything different than the transformers definition. However, this issue persists.
• Swapping to apply the scores to the expert outputs (as is common for most MoEs) does not cause this problem and results in high recovery. As such, enabling this for the time being so that the llama4 pathway does not produce NaN scales
• We can potentially revisit with another dataset but considering how good recovery is, I think this is sufficient to unblock release.
• I have left a note about this deviation from the definition in the modeling code

Evals:

98% Recovery

|   Tasks   |Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----------|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k_llama|      3|flexible_extract|     8|exact_match|↑  |0.934|±  |0.0068|
|           |       |strict_match    |     8|exact_match|↑  |0.931|±  |0.0070|

Greater than 98% Recovery

|      Groups      |Version|   Filter   |n-shot|  Metric   |   |Value |   |Stderr|
|------------------|------:|------------|------|-----------|---|-----:|---|-----:|
|mmlu_llama        |      1|strict_match|      |exact_match|↑  |0.7997|±  |0.0032|
| - humanities     |      1|strict_match|      |exact_match|↑  |0.7696|±  |0.0059|
| - other          |      1|strict_match|      |exact_match|↑  |0.8172|±  |0.0066|
| - social sciences|      1|strict_match|      |exact_match|↑  |0.8781|±  |0.0058|
| - stem           |      0|strict_match|      |exact_match|↑  |0.7510|±  |0.0074|

Greater than 99% recovery

|       Tasks       |Version|   Filter   |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------|------:|------------|-----:|-----------|---|-----:|---|-----:|
|arc_challenge_llama|      1|strict_match|     0|exact_match|↑  |0.9296|±  |0.0075|

Greater than 99% recovery

|  Tasks   |Version|Filter|n-shot|Metric|   |Value |   |Stderr|
|----------|------:|------|-----:|------|---|-----:|---|-----:|
|winogrande|      1|none  |     0|acc   |↑  |0.6835|±  |0.0131|

Greater than 98% recovery

|truthfulqa_mc2|      3|none  |     0|acc        |↑  | 0.6177|±  |0.0164|

[gemini-code-assist]

Summary of Changes

Hello @dsikka, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug in the Llama4 calibration process by refactoring the model's forward pass. The changes aim to improve the accuracy and consistency of expert routing and output accumulation, while also enhancing compatibility with recent versions of the transformers library.

Highlights

• Llama4 Calibration Logic Refactor: The core logic for calculating router scores and accumulating expert outputs within the Llama4 model's forward method has been significantly revised to address calibration issues.
• Transformers Library Compatibility: The change introduces explicit handling for router_logits potentially being a tuple, ensuring compatibility with transformers versions 4.53 and greater.
• Simplified Expert Output Accumulation: The previous conditional logic for calibrate_all_experts and top_token_mask has been replaced with a more direct method of applying router scores and summing expert contributions.
• Debugging Statement Added: A print statement for self.calibrate_all_experts has been added to the __init__ method, likely for debugging purposes during development.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the forward pass of SequentialLlama4TextMoe to fix an issue with MoE calibration. The new implementation correctly handles different transformers library versions and uses a dense pass for calibration. However, it introduces a couple of issues. A debug print statement has been left in the code. More importantly, the calibrate_all_experts flag is now ignored, causing the module to always perform a dense pass and breaking the sparse routing functionality expected when the flag is False. My review includes suggestions to address these points.