Skip to content

Commit

Permalink
Add save_attribute option to quantize_static (#17945)
Browse files Browse the repository at this point in the history
### Description
<!-- Describe your changes. -->
The model with big Constants tensors size: Estimate size of the RWKV
model: ONNX graph (8MB), initializer tensors(200MB), constants (~5.7GB).
The `onnx.save_model` will got error due to the Constants is not output
in external data. Only the initializer tensors are output as external
data. In this change, expose parameter to support the constants in
external data. Model owner can customize the output behavior and still
keep the default behavior.

Quantize the model and output it to local, got issue due to output size
exceed 2GB even set `use_external_data_format=True`. The
`use_external_data_format` flag only outputs initializer tensors to
external data.
Use the falg `convert_attribute` flag to output all tensors to external
data.
```
def convert_model_to_external_data(
    model: ModelProto,
    all_tensors_to_one_file: bool = True,
    location: Optional[str] = None,
    size_threshold: int = 1024,
    include_attribute: bool = False,
) -> None:
    tensors = _get_initializer_tensors(model)
    if include_attribute:
        tensors = _get_all_tensors(model)
        ...
```
The `onnx.external_data_helper.convert_model_to_external_data` support
output the attribute to external with flag `include_attribute=True`.
However, this parameter is hide by the
`onnxruntime\quantization\onnx_model.py` and the constants(`5.7GB)
within the model will got protobuf 2GB limitation issue with default
parameters.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Fix #17944

---------

Co-authored-by: Thiago Crepaldi <[email protected]>
  • Loading branch information
2 people authored and jchen351 committed Oct 18, 2023
1 parent 74145fc commit c48c034
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions onnxruntime/python/tools/quantization/onnx_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,7 @@ def save_model_to_file(self, output_path, use_external_data_format=False):
self.model,
all_tensors_to_one_file=True,
location=Path(output_path).name + ".data",
convert_attribute=True,
)
for init in self.model.graph.initializer:
self._check_init(init, "end")
Expand Down

0 comments on commit c48c034

Please sign in to comment.