[Quantization] Fix per-tensor FP8 MoE parameter identity for EPLB hot-reload#19499
[Quantization] Fix per-tensor FP8 MoE parameter identity for EPLB hot-reload#19499Socratesa wants to merge 2 commits into
Conversation
github-actions
Bot
added
documentation
Summary of ChangesHello, 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 primarily focuses on enhancing the robustness and correctness of FP8 quantization within Mixture-of-Experts (MoE) models, particularly concerning hot-reloading and partial weight updates in distributed environments. It resolves critical issues related to parameter object identity and prevents unintended double-processing of weights, ensuring stable and accurate model operation during dynamic reconfigurations. Additionally, it includes significant refactoring of the MoriEP dispatcher and CUDA graph buffer management for improved efficiency and maintainability. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Ignored Files
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 primarily addresses issues with per-tensor FP8 MoE parameter identity during EPLB hot-reloading, which caused crashes. The core fixes involve using in-place updates for weights and scales to preserve Parameter object identity and its custom attributes, and ensuring that weight post-processing is only applied to the reloaded expert modules. These changes are well-implemented and correctly solve the described problems.
Additionally, the PR includes significant refactoring and cleanup, such as the removal of the mori backend and old JIT kernels, and a reorganization of CUDA graph input buffers. These changes improve code structure and maintainability. I have one suggestion regarding the calculation of skip_token_capacity_threshold in the benchmarking script, which appears to have been simplified incorrectly.
| internal_state = server_info.get("internal_states", [{}]) | ||
| dp_size = internal_state[0].get("dp_size", None) or 1 | ||
| skip_token_capacity_threshold = ( | ||
| internal_state[0].get("memory_usage", {}).get("token_capacity", 1000000000) | ||
| ) |
There was a problem hiding this comment.
The calculation for skip_token_capacity_threshold now only considers the token capacity of the first DP rank. Since batch_size in should_skip_due_to_token_capacity refers to the total batch size, skip_token_capacity_threshold should represent the total capacity across all DP ranks. This change may cause benchmarks to be skipped incorrectly. It would be more accurate to sum the capacities from all DP ranks.
| internal_state = server_info.get("internal_states", [{}]) | |
| dp_size = internal_state[0].get("dp_size", None) or 1 | |
| skip_token_capacity_threshold = ( | |
| internal_state[0].get("memory_usage", {}).get("token_capacity", 1000000000) | |
| ) | |
| internal_state = server_info.get("internal_states", [{}]) | |
| skip_token_capacity_threshold = sum( | |
| state.get("memory_usage", {}).get("token_capacity", 1000000000) | |
| for state in internal_state | |
| ) |
3f98267 to
4f219f8
Compare
The per-tensor FP8 MoE path in process_weights_after_loading() replaced Parameter objects with new ones via torch.nn.Parameter(), destroying custom attributes (weight_loader, quant_method) needed by EPLB weight hot-reload. The block-quant path was already fixed (PR sgl-project#17449) but the per-tensor path was missed. Additionally, update_weights_from_disk() unconditionally called process_weights_after_loading() on ALL modules even during partial reloads (e.g. EPLB expert rebalancing), causing non-expert layers like FP8 attention to be double-processed (double transpose -> shape mismatch). Changes in fp8.py: - Dynamic-quant path: use .data= for weight rebinding and .data[expert].fill_() for scale updates instead of Parameter replacement. - Checkpoint-FP8 path: use .data.fill_() for input_scale merging; fill both columns of the [E,2] w13_weight_scale in-place instead of replacing with a new [E] Parameter. - DeepGemm apply(): add ndim==2 guard to collapse w13_weight_scale from [E,2] to [E] via [:,0] before expanding to block shape. Changes in model_runner.py: - When weight_name_filter is set (EPLB expert rebalancing), split load_weights and process_weights_after_loading into two steps, only calling the latter on modules whose names match the filter. Signed-off-by: Socratesa <lihaode@zju.edu.cn>
4f219f8 to
7ce24c3
Compare
1 participant
Motivation
When using per-tensor FP8 quantized MoE models (e.g.,
RedHatAI/DeepSeek-Coder-V2-Lite-Instruct-FP8) with Elastic EP Load Balancing (EPLB), the failover weight hot-reload crashes with:followed by (after working around the above):
This was also mentioned in #18191.
Issue 1 — Parameter identity loss in per-tensor FP8 MoE
Fp8MoEMethod.process_weights_after_loading()creates newParameterobjects viatorch.nn.Parameter(...)forw13_input_scale,w2_input_scale, andw13_weight_scale. This destroys custom attributes (weight_loader,quant_method) thatdeepseek_weight_loader.pyrelies on during EPLB hot-reload.The block-quant path was already fixed in #17449 via
_copy_or_rebind(), but the per-tensor path was never updated.Issue 2 — Double
process_weights_after_loadingon partial reloadupdate_weights_from_disk()callsload_weights_and_postprocess(), which runsprocess_weights_after_loading()on all quantized modules unconditionally — even when only expert weights were reloaded. For FP8 linear attention layers (e.g.,q_proj), this means the weight gets transposed twice (once during initial load, once during reload), corrupting the shape from[K, N]back to[N, K].Modifications
python/sglang/srt/layers/quantization/fp8.pyDynamic-quant path (non-checkpoint FP8):
.data =for weight rebinding and.data[expert].fill_()for scale updates instead ofParameterreplacement.Checkpoint-FP8 path —
input_scalemerging:torch.nn.Parameter(max_val)withlayer.w13_input_scale.data.fill_(max_val)to preserve the originalParameterobject in-place.Checkpoint-FP8 path —
w13_weight_scalemerging:[E, 2]Parameter with a new[E]one, fill both columns with the same per-expert max in-place. The[E, 2]shape is preserved soweight_loadercan still indexparam.data[expert_id][shard_idx]during reload.apply()— DeepGEMM path:ndim == 2guard to collapsew13_weight_scalefrom[E, 2]to[E]via[:, 0]before expanding to block shape.python/sglang/srt/model_executor/model_runner.pyPartial reload isolation:
weight_name_filteris set (EPLB expert rebalancing), splitload_weightsandprocess_weights_after_loadinginto two steps, only calling the latter on modules whose names match the filter. This prevents double-processing of attention layers and other non-expert modules.Accuracy Tests
Validated with
test_mooncake_ep_small.py(use DeepSeek-Coder-V2-Lite-Instruct-FP8): 7/7 tests passed.Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci