Skip to content
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

Add support for op_block_list #1036

Merged
merged 5 commits into from
Nov 25, 2024
Merged

Add support for op_block_list #1036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b31f40f
Add support for op_block_list
pdufour Nov 18, 2024
ff81cb0
Remove arg
pdufour Nov 18, 2024
2ce32be
Set default to none
pdufour Nov 20, 2024
3e6b01a
Minor code suggestions
xenova Nov 25, 2024
815d6b3
whoops - actually apply suggestions
xenova Nov 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ __pycache__
.vscode
node_modules
.cache
.DS_STORE

# Do not track build artifacts/generated files
/dist
Expand Down
28 changes: 26 additions & 2 deletions scripts/quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from enum import Enum

from tqdm import tqdm
from typing import Set
from typing import Set, List, Optional
import onnx
import os

Expand Down Expand Up @@ -110,6 +110,16 @@ class QuantizationArguments:
},
)

op_block_list: List[str] = field(
default=None,
metadata={
"help": "List of operators to exclude from quantization."
"Can be any standard ONNX operator (see https://onnx.ai/onnx/operators/)"
"or your custom implemented operators.",
"nargs": "+",
},
)


def get_operators(model: onnx.ModelProto) -> Set[str]:
operators = set()
Expand All @@ -131,6 +141,7 @@ def quantize_q8(
per_channel: bool,
reduce_range: bool,
weight_type: QuantType,
op_block_list: Optional[List[str]]
):
"""
Quantize the weights of the model from float32 to int8/uint8
Expand All @@ -140,6 +151,10 @@ def quantize_q8(
it is faster on most CPU architectures
"""

op_types_to_quantize = set(IntegerOpsRegistry.keys())
if op_block_list is not None:
op_types_to_quantize.difference_update(op_block_list)

quantizer = ONNXQuantizer(
model,
per_channel,
Expand All @@ -151,7 +166,7 @@ def quantize_q8(
tensors_range=None,
nodes_to_quantize=[],
nodes_to_exclude=[],
op_types_to_quantize=list(IntegerOpsRegistry.keys()),
op_types_to_quantize=op_types_to_quantize,
extra_options=dict(
EnableSubgraph=True,
MatMulConstBOnly=True,
Expand All @@ -165,6 +180,7 @@ def quantize_q8(
def quantize_fp16(
model: onnx.ModelProto,
save_path: str,
op_block_list: Optional[List[str]]
):
"""
Quantize the weights of the model from float32 to float16
Expand All @@ -174,10 +190,15 @@ def quantize_fp16(
# ValueError: Message onnx.ModelProto exceeds maximum protobuf size of 2GB: 2338583841
disable_shape_infer = model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF

blocked_ops = set(float16.DEFAULT_OP_BLOCK_LIST)
if op_block_list is not None:
blocked_ops.update(op_block_list)

Comment on lines +193 to +196
Copy link
Collaborator

@xenova xenova Nov 25, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

One minor limitation of this updated approach is that you can't choose to quantize a node which is in the default block list. Most of those ops are chosen since there aren't fp16 variants of those ops, so I don't think this is an issue.

TLDR: Can only add to block list.

model_fp16 = float16.convert_float_to_float16(
model,
keep_io_types=True,
disable_shape_infer=disable_shape_infer,
op_block_list=blocked_ops,
)
graph = gs.import_onnx(model_fp16)
graph.toposort()
Expand Down Expand Up @@ -271,6 +292,7 @@ def quantize(input_folder, output_folder, quantization_args: QuantizationArgumen
quantize_fp16(
model,
save_path,
quantization_args.op_block_list
)

elif mode in (QuantMode.Q4, QuantMode.Q4F16):
Expand All @@ -287,6 +309,7 @@ def quantize(input_folder, output_folder, quantization_args: QuantizationArgumen
quantize_fp16(
q4_model,
save_path,
quantization_args.op_block_list,
)

elif mode == QuantMode.BNB4:
Expand Down Expand Up @@ -331,6 +354,7 @@ def quantize(input_folder, output_folder, quantization_args: QuantizationArgumen
per_channel=quantization_args.per_channel,
reduce_range=quantization_args.reduce_range,
weight_type=weight_type,
op_block_list=quantization_args.op_block_list,
)


Expand Down
Loading