[NPU][2/N] Ascend NPU quantization refactoring & more quantization formats support #14504
Conversation
Summary of ChangesHello @OrangeRedeng, 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 significantly refactors and expands the quantization capabilities for Ascend NPU hardware. It introduces a mechanism to automatically detect and apply NPU-optimized quantization settings and adds support for various quantization schemes (W4A4, W4A8, W8A8) for both standard linear layers and Mixture-of-Experts (MoE) layers. The changes aim to leverage NPU's hardware acceleration for improved performance and efficiency in quantized model inference. 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
|
OrangeRedeng
changed the title
![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 a refactoring for Ascend NPU quantization support. It adds several new quantization methods for NPU, including W4A4, W4A8, and W8A8, along with their Mixture-of-Experts (MoE) variants. The changes also include updates to the model configuration to detect and apply these NPU-specific quantization schemes. My review focuses on ensuring correctness, consistency, and code quality. I've identified several critical issues such as missing imports, incorrect class inheritance, and improper use of decorators that could lead to runtime errors. I've also provided suggestions to improve code readability and maintainability. Overall, this is a significant step towards enabling efficient quantization on Ascend NPUs, but the identified issues should be addressed before merging.
| tp_rank: Optional[int] = 0, | ||
| ) -> torch.Tensor: | ||
| original_dtype = x.dtype | ||
| quant_out, dynamic_scale = torch_npu.npu_dynamic_quant( |
There was a problem hiding this comment.
torch_npu is used here but it is not imported in the file. This will raise a NameError at runtime. Please add import torch_npu at the top of the file, guarded by an NPU availability check, similar to other files in this PR.
| class NPU_W4A8DynamicLinearMethod: | ||
| """Linear method for NPU W4A8_DYNAMIC.""" | ||
|
|
||
| def __init__(self): | ||
| self.transpose_weight = True | ||
| try: | ||
| self.group_size = self.quantization_config.get("group_size", 256) | ||
| except AttributeError: | ||
| self.group_size = 256 | ||
|
|
There was a problem hiding this comment.
This class has several issues:
- It should inherit from
_NPULinearMethodBasefor consistency. You will need to import it fromsglang.srt.hardware_backend.npu.quantization.utils. - The
__init__method does not callsuper().__init__(). self.quantization_configis a typo forself.quant_config. This causes thetryblock to always fail, falling back to the defaultgroup_size.
These issues will lead to incorrect behavior. Please apply the suggested fix.
| class NPU_W4A8DynamicLinearMethod: | |
| """Linear method for NPU W4A8_DYNAMIC.""" | |
| def __init__(self): | |
| self.transpose_weight = True | |
| try: | |
| self.group_size = self.quantization_config.get("group_size", 256) | |
| except AttributeError: | |
| self.group_size = 256 | |
| class NPU_W4A8DynamicLinearMethod(_NPULinearMethodBase): | |
| """Linear method for NPU W4A8_DYNAMIC.""" | |
| def __init__(self, quant_config: Optional["QuantizationConfig"] = None): | |
| super().__init__(quant_config) | |
| self.transpose_weight = True | |
| if self.quant_config: | |
| self.group_size = self.quant_config.get("group_size", 256) | |
| else: | |
| self.group_size = 256 | |
| def __init__(self): | ||
| self.transpose_weight = True | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
The apply method is decorated with @staticmethod but its first argument is self. This is incorrect and will lead to unexpected behavior or errors at runtime because the self argument will be incorrectly bound. Based on its usage, it should be an instance method. Please remove the @staticmethod decorator.
| def __init__(self): | ||
| self.transpose_weight = True | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
The apply method is decorated with @staticmethod but its first argument is self. This is incorrect and will lead to unexpected behavior or errors at runtime because the self argument will be incorrectly bound. Based on its usage, it should be an instance method. Please remove the @staticmethod decorator.
| from sglang.srt.hardware_backend.npu.utils import npu_format_cast | ||
| from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase | ||
|
|
||
| class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): |
There was a problem hiding this comment.
The filename w4a16_moe.py is inconsistent with the class name NPUW4A8Int4DynamicMoEMethod it contains. The class name suggests W4A8 quantization, while the filename suggests W4A16. To avoid confusion, please rename the file to align with its content, for example, to w4a8_moe.py.
| class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): | ||
| """Linear method for NPU W8A8.""" | ||
|
|
||
| def __init__(self): | ||
| self.transpose_weight = True |
There was a problem hiding this comment.
The __init__ method of NPUW8A8Int8LinearMethod should call super().__init__() to ensure the base class is properly initialized. It's also good practice to accept quant_config to be consistent with the base class signature.
| class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): | |
| """Linear method for NPU W8A8.""" | |
| def __init__(self): | |
| self.transpose_weight = True | |
| class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): | |
| """Linear method for NPU W8A8.""" | |
| def __init__(self, quant_config: Optional["QuantizationConfig"] = None): | |
| super().__init__(quant_config) | |
| self.transpose_weight = True |
| class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): | ||
| """Linear method for NPU W8A8_DYNAMIC.""" | ||
|
|
||
| def __init__(self): | ||
| self.transpose_weight = True |
There was a problem hiding this comment.
The __init__ method of NPUW8A8Int8DynamicLinearMethod should call super().__init__() to ensure the base class is properly initialized. It's also good practice to accept quant_config to be consistent with the base class signature.
| class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): | |
| """Linear method for NPU W8A8_DYNAMIC.""" | |
| def __init__(self): | |
| self.transpose_weight = True | |
| class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): | |
| """Linear method for NPU W8A8_DYNAMIC.""" | |
| def __init__(self, quant_config: Optional["QuantizationConfig"] = None): | |
| super().__init__(quant_config) | |
| self.transpose_weight = True |
| def _find_quant_modelslim_config(self): | ||
| quant_config_file = Path(self.model_path, "quant_model_description.json") | ||
| if quant_config_file.is_file(): | ||
| with open(quant_config_file) as f: | ||
| quant_cfg = json.load(f) | ||
| else: | ||
| quant_cfg = None | ||
|
|
||
| return quant_cfg |
There was a problem hiding this comment.
This function can be simplified by returning early if the file exists. This would make the code more concise and readable by reducing nesting.
| unpacked_w13_weight = ( | ||
| self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4) | ||
| .view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1) | ||
| .transpose(1, 2) | ||
| .contiguous() | ||
| .int() | ||
| ) | ||
| unpacked_w2_weight = ( | ||
| self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4) | ||
| .view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1) | ||
| .transpose(1, 2) | ||
| .contiguous() | ||
| .int() | ||
| ) |
There was a problem hiding this comment.
These long chained method calls are difficult to read and debug. Please consider breaking them into multiple lines using intermediate variables for better clarity. This applies to both unpacked_w13_weight and unpacked_w2_weight.
| sscale_uint64_tensor = sscale_uint64_tensor.npu() | ||
| return sscale_uint64_tensor, bias | ||
|
|
||
| def update_bias(self, layer, w13_bias, w2_bias): |
There was a problem hiding this comment.
The parameters w13_bias and w2_bias are unused in the update_bias method. They are always None when passed from process_weights_after_loading. Please remove them from the method signature to simplify the code and update the call at line 87 to self.update_bias(layer).
| def update_bias(self, layer, w13_bias, w2_bias): | |
| def update_bias(self, layer): |
1 participant
Motivation
DRAFT !!!
Related to #14424
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist