Skip to content

fix: llama 4 + EP IndexError #12636

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: llama 4 + EP IndexError #12636

wants to merge 1 commit into from

Conversation

@b8zhong
Copy link
Collaborator

@b8zhong b8zhong commented Nov 4, 2025 •
edited
Loading

It seems Llama 4 EP previously did not work.

python3 -m sglang.launch_server \
  --model-path=/opt/dlami/nvme/models/Llama-4-Scout-17B-16E-Instruct/ \
  --tp=8 \
  --ep-size=8 \
  --trust-remote-code \
  --mem-fraction-static=0.7 \
  --context-length=131072 \
  --kv-cache-dtype=fp8_e4m3 \
  --attention-backend=fa3 \
  --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 8}'

           ^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/model_loader/loader.py", line 599, in load_model
    self.load_weights_and_postprocess(
  File "/sgl-workspace/sglang/python/sglang/srt/model_loader/loader.py", line 607, in load_weights_and_postprocess
    model.load_weights(weights)
  File "/sgl-workspace/sglang/python/sglang/srt/models/mllama4.py", line 675, in load_weights
    if self._handle_expert_weights(
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/mllama4.py", line 768, in _handle_expert_weights
    return self._handle_expert_weight_params(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/mllama4.py", line 934, in _handle_expert_weight_params
    weight_loader(
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py", line 500, in weight_loader
    self._weight_loader_impl(
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py", line 608, in _weight_loader_impl
    expert_data = param.data[expert_id]
                  ~~~~~~~~~~^^^^^^^^^^^
IndexError: index 2 is out of bounds for dimension 0 with size 2

[2025-11-04 01:13:29] Received sigquit from a child process. It usually means the child failed.
Killed

Fix #12577

My understanding of EP is naive, but here is my explanation:

When expert-location metadata isn’t available, the loader was indexing param.data[expert_id] with a global ID (e.g., 2) while the param only has num_local_experts rows (2), causing the IndexError.
To avoid breaking other models (DSR1) I think CI should be able to catch if it does.

Loading is fixed + accuracy is fine

root@ip-10-40-12-14:/sgl-workspace/sglang# python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 200
Downloading from https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl to /tmp/test.jsonl
/tmp/test.jsonl: 732kB [00:00, 23.3MB/s]                                                                                                                                                                                                                                     
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [01:00<00:00, 21.91it/s]
Accuracy: 0.914
Invalid: 0.000
Latency: 60.419 s
Output throughput: 2238.073 token/s

@b8zhong b8zhong requested a review from Copilot November 4, 2025 17:40
@gemini-code-assist
Copy link
Contributor

gemini-code-assist bot commented Nov 4, 2025

Summary of Changes

Hello @b8zhong, I'm Gemini Code Assist1! 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 critical bug that caused Llama 4 models to fail loading with Expert Parallelism due to an IndexError in the weight loading process. The change specifically targets how expert weights are handled when global metadata is absent, ensuring that expert IDs are correctly interpreted and mapped. This fix enables robust model initialization for Llama 4 with EP, with verified preservation of model accuracy.

Highlights

  • Llama 4 Expert Parallelism Fix: Resolves an IndexError that previously prevented Llama 4 models from loading correctly when utilizing Expert Parallelism (EP).
  • Correct Expert Weight Loading: Modifies the weight loading mechanism to correctly handle expert IDs when global expert location metadata is unavailable, ensuring proper mapping of global expert IDs to local ones.
  • Accuracy Verification: Confirms that the implemented fix maintains model accuracy, as validated by benchmark results on the GSM8K dataset.
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. ↩

@b8zhong b8zhong changed the title fix: llama 4 + EP fix: llama 4 + EP IndexError Nov 4, 2025
@b8zhong b8zhong added the run-ci label Nov 4, 2025
gemini-code-assist[bot]
gemini-code-assist bot reviewed Nov 4, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request addresses a critical IndexError that occurred during weight loading for Llama 4 models with Expert Parallelism (EP). The root cause was the incorrect use of a global expert ID to index a parameter array expecting a local expert ID when expert location metadata was not available. The fix correctly redirects the weight loading process to the _weight_loader_physical method, which properly handles the global-to-local expert ID mapping. This change is a direct and effective solution to the bug, and the provided benchmark results confirm its correctness. The code change is sound.

Copilot AI reviewed Nov 4, 2025
View reviewed changes
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR fixes the weight loading logic for MoE (Mixture of Experts) layers when global expert location metadata is not available. The key change ensures that expert ID mapping from global to local is correctly handled in all cases.

  • Changed the method call from _weight_loader_impl to _weight_loader_physical when global expert location metadata is absent
  • Added explanatory comment clarifying that the physical loader handles global-to-local expert ID mapping
  • Ensures consistent expert ID handling regardless of whether global metadata is present

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner

@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner

@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner

@kushanam kushanam Awaiting requested review from kushanam kushanam is a code owner

@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

run-ci

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug] IndexError with Llama 4 + EP

1 participant

@b8zhong