Skip to content

Commit

Permalink
Fix weight loading of large model in INC quantzation pass (#857)
Browse files Browse the repository at this point in the history
## Describe your changes

Fix weight loading of large model in inc quantzation pass. Reload weight
for model > 2GB to prevent missing weight files.

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [ ] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link

#852

---------

Signed-off-by: yuwenzho <[email protected]>
  • Loading branch information
yuwenzho committed Jan 4, 2024
1 parent 0281211 commit 7b8a267
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions olive/passes/onnx/inc_quantization.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,6 +562,13 @@ def _run_for_config(
"find any quantized model which meet accuracy goal. "
"Try to increase 'max_trials' in 'tuning_criterion'."
)

# reload weight for model with size > 2GB to prevent error of missing weight files
if q_model.is_large_model:
from onnx.external_data_helper import load_external_data_for_model

load_external_data_for_model(q_model.model, os.path.dirname(q_model._model_path))

# save the model to the output path and return the model
return model_proto_to_olive_model(q_model.model, output_model_path, config)

Expand Down

0 comments on commit 7b8a267

Please sign in to comment.