Skip to content

Commit

Permalink
Add save_attribute option to quantize_static (microsoft#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 microsoft#17944

---------

Co-authored-by: Thiago Crepaldi <[email protected]>
  • Loading branch information
zhipenghan and Thiago Crepaldi authored Oct 14, 2023
1 parent 11af344 commit a55b268
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 a55b268

Please sign in to comment.