Skip to content

Fix onnx calibration state restoration - #4115

Open
Sh0rck_Wang (WangXukang-cypher) wants to merge 2 commits into
qualcomm:developfrom
WangXukang-cypher:fix-onnx-calibration-state-restoration
Open

Fix onnx calibration state restoration#4115
Sh0rck_Wang (WangXukang-cypher) wants to merge 2 commits into
qualcomm:developfrom
WangXukang-cypher:fix-onnx-calibration-state-restoration

Conversation

@WangXukang-cypher

Copy link
Copy Markdown

Fixes #4114

Summary

The ONNX compute_encodings context previously left quantizers in a partially calibrated state if the calibration callback or encoding finalization raised an exception.

This change makes calibration exception-safe by:

  • Preserving enabled quantizers' existing encodings and operation modes before calibration.
  • Restoring the preserved encodings and operation modes when calibration exits with an exception.
  • Re-raising the original exception after restoring quantizer state.
  • Making nested compute_encodings contexts reuse the outer calibration lifecycle so that an inner context does not reset statistics or finalize encodings prematurely.
  • Adding regression tests for exceptional exits and nested calibration contexts.

The exceptional-exit regression test verifies that:

  • Existing encoding values are preserved.
  • Quantizer operation modes are restored.
  • Quantized forward outputs remain unchanged after a failed calibration attempt.

Test plan

The changes were tested in a Python 3.10 CPU environment with an editable source build of AIMET ONNX.

  • Run the new exceptional-exit and nested-context regression tests.
  • Run the existing normal compute_encodings tests.
  • Run all compute_encodings-related tests in test_quantsim.py.
  • Run repository pre-commit checks on the modified files.
  • Run IDE lint checks on the modified files.

Results

4 passed

Signed-off-by: Sh0rck_Wang <wangxukang010406@outlook.com>
Signed-off-by: Sh0rck_Wang <wangxukang010406@outlook.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ONNX calibration does not restore quantizer state after an exception

1 participant