Dev quant #421
Dev quant #421
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @gushiqiao, 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 enhances the quantization capabilities of the system by integrating VLLM's efficient quantization operations for both FP8 and INT8. It introduces a flexible mechanism to configure block-wise quantization for FP8 models and updates existing Deepseekv3 configurations to utilize these new features, alongside general code clean-up.
Highlights
- New Quantization Modules: Added
VllmQuantLinearInt8andVllmQuantLinearFp8classes tomodule_utils.pyto leverage VLLM's custom operations for efficient INT8 and FP8 linear layer quantization, respectively. - Configurable Block-wise FP8 Quantization: Introduced a
block_wise_quantflag in model configurations andBaseModelto allow switching between the existing LLMC FP8 block-wise quantization and the new VLLM FP8 quantization. - Expanded Quantization Support: The
BaseModelnow dynamically selects the appropriate quantization linear layer (LLMC'sLlmcFp8Linearor VLLM'sVllmQuantLinearFp8/VllmQuantLinearInt8) based on the specifiedtorch_dtypeandblock_wise_quantsettings, extending support to INT8. - Deepseekv3 Configuration Updates: Updated several Deepseekv3 quantization configuration files (
awq_w_only_dsv3.yml,osplus_w_a_dsv3.yml,quarot_w_a_dsv3.yml,rtn_w_a_dsv3.yml,rtn_w_only_dsv3.yml,smoothquant_w_a_dsv3.yml) to enable the newblock_wise_quantsetting. - Code Clean-up and Refactoring: Removed unused imports (
FloatQuantizer,calculate_zeros_width,partial) and refactored the weight loading function name fromload_fp8_weighttoload_quant_weightfor broader applicability.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 is currently in preview and 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 to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces support for different quantization schemes. The changes are mostly centered around llmc/models/base_model.py and llmc/compression/quantization/module_utils.py. I've identified issues related to batch processing and a potential runtime error that need to be addressed.
| class VllmQuantLinearInt8(nn.Module): | ||
| def __init__(self, in_features, out_features, bias=True): | ||
| super().__init__() | ||
| self.in_features = in_features | ||
| self.out_features = out_features | ||
|
|
||
| self.register_buffer('weight', torch.empty((out_features, in_features), dtype=torch.int8)) | ||
| self.register_buffer('weight_scale', torch.empty((out_features, 1), dtype=torch.float32)) | ||
|
|
||
| if bias: | ||
| self.register_buffer('bias', torch.empty(out_features, dtype=torch.bfloat16)) | ||
| else: | ||
| self.register_buffer('bias', None) | ||
|
|
||
| def act_quant_func(self, x): | ||
| input_tensor_quant, input_tensor_scale, _ \ | ||
| = ops.scaled_int8_quant(x, scale=None, azp=None, symmetric=True) | ||
| return input_tensor_quant, input_tensor_scale | ||
|
|
||
| def forward(self, input_tensor): | ||
| input_tensor = input_tensor.squeeze(0) | ||
| shape = (input_tensor.shape[0], self.weight.shape[0]) | ||
| dtype = input_tensor.dtype | ||
| device = input_tensor.device | ||
| output_tensor = torch.empty(shape, dtype=dtype, device=device, requires_grad=False) | ||
|
|
||
| input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) | ||
| torch.ops._C.cutlass_scaled_mm( | ||
| output_tensor, | ||
| input_tensor_quant, | ||
| self.weight.t(), | ||
| input_tensor_scale, | ||
| self.weight_scale.float(), | ||
| self.bias, | ||
| ) | ||
| return output_tensor.unsqueeze(0) | ||
|
|
||
| @classmethod | ||
| @torch.no_grad() | ||
| def new(cls, module): | ||
| in_features = module.in_features | ||
| out_features = module.out_features | ||
| bias = module.bias is not None | ||
| new_module = cls(in_features, out_features, bias) | ||
| return new_module |
There was a problem hiding this comment.
The forward method in VllmQuantLinearInt8 assumes a batch size of 1 due to the use of input_tensor.squeeze(0) and output_tensor.unsqueeze(0). This will cause issues when processing batches of data with a size greater than 1. Reshape the input tensor to handle arbitrary batch sizes.
def forward(self, input_tensor):
input_shape = input_tensor.shape
input_tensor = input_tensor.view(-1, self.in_features)
output_tensor = torch.empty(
(input_tensor.shape[0], self.out_features),
dtype=input_tensor.dtype,
device=input_tensor.device,
requires_grad=False,
)
input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor)
torch.ops._C.cutlass_scaled_mm(
output_tensor,
input_tensor_quant,
self.weight.t(),
input_tensor_scale,
self.weight_scale.float(),
self.bias,
)
return output_tensor.view(*input_shape[:-1], self.out_features)| class VllmQuantLinearFp8(nn.Module): | ||
| def __init__(self, in_features, out_features, bias=True): | ||
| super().__init__() | ||
| self.in_features = in_features | ||
| self.out_features = out_features | ||
| self.register_buffer('weight', torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn)) # noqa | ||
| self.register_buffer('weight_scale', torch.empty((out_features, 1), dtype=torch.float32)) | ||
| if bias: | ||
| self.register_buffer('bias', torch.empty(out_features, dtype=torch.bfloat16)) | ||
| else: | ||
| self.register_buffer('bias', None) | ||
|
|
||
| def act_quant_func(self, x): | ||
| input_tensor_quant, input_tensor_scale \ | ||
| = ops.scaled_fp8_quant(x, None, scale_ub=None, use_per_token_if_dynamic=True) | ||
| return input_tensor_quant, input_tensor_scale | ||
|
|
||
| def forward(self, input_tensor): | ||
| input_tensor = input_tensor.squeeze(0) | ||
| shape = (input_tensor.shape[0], self.weight.shape[0]) | ||
| dtype = input_tensor.dtype | ||
| device = input_tensor.device | ||
| output_tensor = torch.empty(shape, dtype=dtype, device=device, requires_grad=False) | ||
| input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor) | ||
| torch.ops._C.cutlass_scaled_mm( | ||
| output_tensor, | ||
| input_tensor_quant, | ||
| self.weight.t(), | ||
| input_tensor_scale, | ||
| self.weight_scale.float(), | ||
| self.bias, | ||
| ) | ||
|
|
||
| return output_tensor.unsqueeze(0) | ||
|
|
||
| @classmethod | ||
| @torch.no_grad() | ||
| def new(cls, module): | ||
| in_features = module.in_features | ||
| out_features = module.out_features | ||
| bias = module.bias is not None | ||
| new_module = cls(in_features, out_features, bias) | ||
| return new_module |
There was a problem hiding this comment.
The forward method in VllmQuantLinearFp8 also assumes a batch size of 1 due to the use of input_tensor.squeeze(0) and output_tensor.unsqueeze(0). This needs to be corrected to properly support batched inputs by reshaping the input tensor.
def forward(self, input_tensor):
input_shape = input_tensor.shape
input_tensor = input_tensor.view(-1, self.in_features)
output_tensor = torch.empty(
(input_tensor.shape[0], self.out_features),
dtype=input_tensor.dtype,
device=input_tensor.device,
requires_grad=False,
)
input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor)
torch.ops._C.cutlass_scaled_mm(
output_tensor,
input_tensor_quant,
self.weight.t(),
input_tensor_scale,
self.weight_scale.float(),
self.bias,
)
return output_tensor.view(*input_shape[:-1], self.out_features)| self.load_fp8_weight() | ||
| params_dict) | ||
|
|
||
| self.load_quant_weight() | ||
|
|
||
| logger.info(f'fp8 block size: {self.fp8_block_size}') |
There was a problem hiding this comment.
The logger.info call for fp8 block size can lead to an AttributeError if self.torch_dtype is not torch.float8_e4m3fn or if self.block_wise_quant is false, because self.fp8_block_size is only defined within that conditional block. Move the log statement inside the if self.block_wise_quant: block.
| logger.info(f'fp8 block size: {self.fp8_block_size}') | |
| # logger.info(f'fp8 block size: {self.fp8_block_size}') |
No description provided.