fix: llama 4 fp4 loading #12649
fix: llama 4 fp4 loading #12649
Conversation
b8zhong
requested review from
BBuf,
Edwardf0t1,
FlamingoPg,
HaiShaw,
Ying1123,
ch-wan,
ispobock,
kushanam,
merrymercy and
zhyncs
as code owners
Summary of ChangesHello @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 critical issues related to loading Llama 4 models, especially when utilizing FP4 quantization and Mixture-of-Experts (MoE) layers. The primary goal is to ensure accurate and efficient loading of quantized weights and scales, which was previously leading to suboptimal performance as indicated by the provided benchmarks. The changes introduce specialized handling for MoE layer components, improve integration with Flashinfer kernels, and refine weight permutation logic for rotary embeddings, ultimately aiming to stabilize and enhance the model's performance with FP4 quantization. Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces fixes for loading Llama 4 FP4 models, which appears to address performance issues observed in benchmarks. The core changes involve a significant refactoring of the weight loading logic in llama4.py and mllama4.py to correctly handle permutations and formats of expert and attention weights. Additionally, there are improvements in handling CUDA graph capture with empty inputs and aligning data types for FlashInfer kernels. My review focuses on improving code maintainability by removing redundant code blocks and adhering to Python's standard style guidelines for imports. The core logic of the fix seems sound and well-targeted.
| num_experts: int, | ||
| loaded_params: set, | ||
| ) -> bool: | ||
| import re |
There was a problem hiding this comment.
The import re statement is located inside the _handle_expert_scale_params method. According to the PEP 8 style guide, imports should be placed at the top of the file. Moving this import to the beginning of python/sglang/srt/models/llama4.py will improve code style and consistency.
| else: | ||
| for expert_id in range(num_experts): | ||
| weight_loader( | ||
| param, | ||
| weight_chunk[expert_id], | ||
| param_name, | ||
| shard_id=shard_id, | ||
| expert_id=expert_id, | ||
| ) |
There was a problem hiding this comment.
The else block here is identical to the elif block above it. This redundancy can make the code harder to maintain. If the logic is indeed the same, the else block can be removed to simplify the code.
1 participant
Don't merge this yet
Rebase #9526
python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 500It's too low, MMLU Pro needs to be >=0.74. It seems to be some accuracy issue in flashinfer_cutlass MoE usage
BF16 launch commands (H200):