Key Error auto_quant.optimize

[Eman-Su]

Hello,
I'm trying to quantize a transformer model using AutoQuant. When I ran the following command:
model, optimized_accuracy, encoding_path = auto_quant.optimize(allowed_accuracy_drop=0.01)
I got the following:
`--------------------------------------------------------------------------
KeyError Traceback (most recent call last)
Cell In[23], line 3
1 # Step 7. Run AutoQuant
2 sim, initial_accuracy = auto_quant.run_inference()
----> 3 model, optimized_accuracy, encoding_path = auto_quant.optimize(allowed_accuracy_drop=0.01)

File ~/.local/lib/python3.10/site-packages/aimet_torch/auto_quant_v2.py:344, in AutoQuant.optimize(self, allowed_accuracy_drop)
337 def optimize(self, allowed_accuracy_drop: float = 0.0) -> Tuple[torch.nn.Module, float, str]:
338 """
339 Integrate and apply post-training quantization techniques.
340
341 :param allowed_accuracy_drop: Maximum allowed accuracy drop
342 :return: Tuple of (best model, eval score, encoding path)
343 """
--> 344 result = self._optimize_helper(self._optimize_main, allowed_accuracy_drop)
345 return result["model"],
346 result["accuracy"],
347 result["encoding_path"]

File ~/.local/lib/python3.10/site-packages/aimet_torch/auto_quant_v2.py:603, in AutoQuant._optimize_helper(self, optimize_fn, allowed_accuracy_drop)
600 _logger.info("Target eval score: %f", target_acc)
601 _logger.info("FP32 eval score (W32A32): %f", self._fp32_acc)
--> 603 ret = optimize_fn(self.fp32_model, target_acc)
605 acc = ret["accuracy"]
606 if acc is not None:

File ~/.local/lib/python3.10/site-packages/aimet_torch/auto_quant_v2.py:728, in AutoQuant._optimize_main(self, fp32_model, target_acc)
726 # Batchnorm Folding
727 with self.eval_manager.session("Batchnorm Folding", ptq=True) as sess:
--> 728 model, _ = sess.wrap(self._apply_batchnorm_folding)(fp32_model)
729 if sess.ptq_result is None:
730 sess.set_ptq_result(model=model,
731 applied_techniques=["batchnorm_folding"],
732 export_kwargs=self._export_kwargs)

File ~/.local/lib/python3.10/site-packages/aimet_torch/auto_quant_v2.py:1032, in _EvalSession.wrap..wrapper(*args, **kwargs)
1029 @functools.wraps(fn)
1030 def wrapper(*args, **kwargs):
1031 if self._cached_result:
-> 1032 return self._cached_result.load()
1033 ret = fn(*args, **kwargs)
1034 self._cached_result = CachedResult(ret)

File ~/.local/lib/python3.10/site-packages/aimet_torch/auto_quant_v2.py:1027, in _EvalSession.wrap..CachedResult.load(self)
1025 """Load cached result """
1026 with open(self._filename, "rb") as f:
-> 1027 return pickle.load(f)

File ~/.local/lib/python3.10/site-packages/torch/fx/graph_module.py:105, in reduce_graph_module(body, import_block)
103 fn_src = body.get('_code') or body['code']
104 forward = _forward_from_src(import_block + fn_src, {})
--> 105 return _deserialize_graph_module(forward, body)

File ~/.local/lib/python3.10/site-packages/torch/fx/graph_module.py:164, in _deserialize_graph_module(forward, body)
161 com = CodeOnlyModule(body)
163 tracer_extras = body.get('_tracer_extras', {})
--> 164 graph = KeepModules().trace(com, **tracer_extras)
166 # Manually set Tracer class on the reconstructed Graph, to avoid
167 # referencing the private local subclass KeepModules.
168 graph._tracer_cls = tracer_cls

File ~/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:730, in Tracer.trace(self, root, concrete_args)
721 patcher.patch_method(
722 torch.nn.Module,
723 "getattr",
724 module_getattr_wrapper,
725 deduplicate=False,
726 )
727 patcher.patch_method(
728 torch.nn.Module, "call", module_call_wrapper, deduplicate=False
729 )
--> 730 _patch_wrapped_functions(patcher)
731 _autowrap_check(patcher, fn_globals, self._autowrap_function_ids)
732 for module in self._autowrap_search:

File ~/.local/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py:931, in _patch_wrapped_functions(patcher)
929 orig_fn = getattr(builtins, name)
930 else:
--> 931 orig_fn = frame_dict[name]
932 patcher.patch(frame_dict, name, _create_wrapped_func(orig_fn))
934 for cls, name in _wrapped_methods_to_patch:

KeyError: 'torch.fx._symbolic_trace._assert_is_none'`

Any pointers on how to solve this ?

Also, when I ran this command:
sim, initial_accuracy = auto_quant.run_inference()

I received the following:
`2024-05-29 11:12:43,664 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.module_sub}
2024-05-29 11:12:43,667 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.module_mul}
2024-05-29 11:12:43,669 - ModelPreparer - INFO - Functional : Adding new module for node: {module_add}
2024-05-29 11:12:43,670 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.module_add_1}
2024-05-29 11:12:43,672 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.c_attn.module_add_2}
2024-05-29 11:12:43,674 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.c_attn.module_addmm}
2024-05-29 11:12:43,675 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_floordiv}
2024-05-29 11:12:43,679 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_floordiv_1}
2024-05-29 11:12:43,681 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_floordiv_2}
2024-05-29 11:12:43,683 - ModelPreparer - INFO - Functional : Adding new module for node: {module_matmul}
2024-05-29 11:12:43,684 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_pow_1}
2024-05-29 11:12:43,686 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_truediv}
2024-05-29 11:12:43,689 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_where}
2024-05-29 11:12:43,692 - ModelPreparer - INFO - Functional : Adding new module for node: {transformer.h.0.attn.module_add_3}
2024-05-29 11:12:43,694 - ModelPreparer - INFO - Functional : Adding new module for node: {module_matmul_1}
.
.
.

2024-05-29 11:12:46,368 - Utils - INFO - Running validator check
2024-05-29 11:12:46,466 - Utils - INFO - Running validator check
2024-05-29 11:19:18,733 - Utils - INFO - All validation checks passed.
2024-05-29 11:19:18,736 - AutoQuant - INFO - Model validation has succeeded. Proceeding to AutoQuant algorithm.
Batchnorm Folding

2024-05-29 11:26:04,784 - BatchNormFolding - INFO - 0 BatchNorms' weights got converted
2024-05-29 11:32:40,154 - Quant - INFO - No config file provided, defaulting to config file at /home/user/.local/lib/python3.10/site-packages/aimet_common/quantsim_config/default_config.json
2024-05-29 11:32:40,205 - Quant - INFO - Unsupported op type Squeeze
2024-05-29 11:32:40,209 - Quant - INFO - Unsupported op type Mean
2024-05-29 11:32:43,698 - Quant - INFO - Selecting DefaultOpInstanceConfigGenerator to compute the specialized config. hw_version:default
2024-05-29 12:42:31,157 - Quant - WARNING - Exporting encodings to yaml will be deprecated in a future release. Ensure that your code can work with the exported files ending in ".encodings" which are saved using json format. For the time being, if yaml export is needed, set aimet_common.utils.SAVE_TO_YAML to True.
2024-05-29 12:51:03,344 - Utils - WARNING - end-marker seen without passing start-marker for 'transformer.h.0.attn.c_attn.module_add_2', continue to use parent context ''
2024-05-29 12:51:03,346 - Utils - WARNING - end-marker seen without passing start-marker for 'transformer.h.0.attn.c_attn.module_add_2', continue to use parent context ''
2024-05-29 12:51:03,350 - Utils - WARNING - end-marker seen without passing start-marker for 'transformer.h.0.attn.c_proj.module_add_4', continue to use parent context ''
.
.
.
.

• Quant - ERROR - Param tensor {transformer.h.0.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,163 - Quant - ERROR - Param tensor {transformer.h.1.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,176 - Quant - ERROR - Param tensor {transformer.h.1.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,184 - Quant - ERROR - Param tensor {transformer.h.2.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,211 - Quant - ERROR - Param tensor {transformer.h.2.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,221 - Quant - ERROR - Param tensor {transformer.h.3.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,234 - Quant - ERROR - Param tensor {transformer.h.3.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,243 - Quant - ERROR - Param tensor {transformer.h.4.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,256 - Quant - ERROR - Param tensor {transformer.h.4.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,264 - Quant - ERROR - Param tensor {transformer.h.5.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,277 - Quant - ERROR - Param tensor {transformer.h.5.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,286 - Quant - ERROR - Param tensor {transformer.h.6.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,310 - Quant - ERROR - Param tensor {transformer.h.6.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,320 - Quant - ERROR - Param tensor {transformer.h.7.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,332 - Quant - ERROR - Param tensor {transformer.h.7.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,341 - Quant - ERROR - Param tensor {transformer.h.8.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,354 - Quant - ERROR - Param tensor {transformer.h.8.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,365 - Quant - ERROR - Param tensor {transformer.h.9.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,380 - Quant - ERROR - Param tensor {transformer.h.9.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,389 - Quant - ERROR - Param tensor {transformer.h.10.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,402 - Quant - ERROR - Param tensor {transformer.h.10.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,411 - Quant - ERROR - Param tensor {transformer.h.11.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,423 - Quant - ERROR - Param tensor {transformer.h.11.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,432 - Quant - ERROR - Param tensor {transformer.h.12.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,445 - Quant - ERROR - Param tensor {transformer.h.12.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,454 - Quant - ERROR - Param tensor {transformer.h.13.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,469 - Quant - ERROR - Param tensor {transformer.h.13.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,479 - Quant - ERROR - Param tensor {transformer.h.14.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,492 - Quant - ERROR - Param tensor {transformer.h.14.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,501 - Quant - ERROR - Param tensor {transformer.h.15.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,513 - Quant - ERROR - Param tensor {transformer.h.15.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,522 - Quant - ERROR - Param tensor {transformer.h.16.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,535 - Quant - ERROR - Param tensor {transformer.h.16.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,544 - Quant - ERROR - Param tensor {transformer.h.17.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,557 - Quant - ERROR - Param tensor {transformer.h.17.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,566 - Quant - ERROR - Param tensor {transformer.h.18.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,579 - Quant - ERROR - Param tensor {transformer.h.18.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,612 - Quant - ERROR - Param tensor {transformer.h.19.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,626 - Quant - ERROR - Param tensor {transformer.h.19.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,635 - Quant - ERROR - Param tensor {transformer.h.20.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,648 - Quant - ERROR - Param tensor {transformer.h.20.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,657 - Quant - ERROR - Param tensor {transformer.h.21.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,670 - Quant - ERROR - Param tensor {transformer.h.21.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,679 - Quant - ERROR - Param tensor {transformer.h.22.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,710 - Quant - ERROR - Param tensor {transformer.h.22.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,720 - Quant - ERROR - Param tensor {transformer.h.23.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,733 - Quant - ERROR - Param tensor {transformer.h.23.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,743 - Quant - ERROR - Param tensor {transformer.h.24.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,758 - Quant - ERROR - Param tensor {transformer.h.24.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,767 - Quant - ERROR - Param tensor {transformer.h.25.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,780 - Quant - ERROR - Param tensor {transformer.h.25.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,789 - Quant - ERROR - Param tensor {transformer.h.26.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,823 - Quant - ERROR - Param tensor {transformer.h.26.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,832 - Quant - ERROR - Param tensor {transformer.h.27.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,844 - Quant - ERROR - Param tensor {transformer.h.27.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,854 - Quant - ERROR - Param tensor {transformer.h.28.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,866 - Quant - ERROR - Param tensor {transformer.h.28.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,875 - Quant - ERROR - Param tensor {transformer.h.29.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,887 - Quant - ERROR - Param tensor {transformer.h.29.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,918 - Quant - ERROR - Param tensor {transformer.h.30.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,933 - Quant - ERROR - Param tensor {transformer.h.30.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,942 - Quant - ERROR - Param tensor {transformer.h.31.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,954 - Quant - ERROR - Param tensor {transformer.h.31.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,963 - Quant - ERROR - Param tensor {transformer.h.32.ln_1.weight} not found in valid param set 2024-05-29 12:51:30,976 - Quant - ERROR - Param tensor {transformer.h.32.ln_2.weight} not found in valid param set 2024-05-29 12:51:30,984 - Quant - ERROR - Param tensor {transformer.h.33.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,009 - Quant - ERROR - Param tensor {transformer.h.33.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,019 - Quant - ERROR - Param tensor {transformer.h.34.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,031 - Quant - ERROR - Param tensor {transformer.h.34.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,040 - Quant - ERROR - Param tensor {transformer.h.35.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,052 - Quant - ERROR - Param tensor {transformer.h.35.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,061 - Quant - ERROR - Param tensor {transformer.h.36.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,073 - Quant - ERROR - Param tensor {transformer.h.36.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,082 - Quant - ERROR - Param tensor {transformer.h.37.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,111 - Quant - ERROR - Param tensor {transformer.h.37.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,121 - Quant - ERROR - Param tensor {transformer.h.38.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,134 - Quant - ERROR - Param tensor {transformer.h.38.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,142 - Quant - ERROR - Param tensor {transformer.h.39.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,155 - Quant - ERROR - Param tensor {transformer.h.39.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,164 - Quant - ERROR - Param tensor {transformer.h.40.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,181 - Quant - ERROR - Param tensor {transformer.h.40.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,208 - Quant - ERROR - Param tensor {transformer.h.41.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,229 - Quant - ERROR - Param tensor {transformer.h.41.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,238 - Quant - ERROR - Param tensor {transformer.h.42.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,251 - Quant - ERROR - Param tensor {transformer.h.42.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,259 - Quant - ERROR - Param tensor {transformer.h.43.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,272 - Quant - ERROR - Param tensor {transformer.h.43.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,280 - Quant - ERROR - Param tensor {transformer.h.44.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,293 - Quant - ERROR - Param tensor {transformer.h.44.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,302 - Quant - ERROR - Param tensor {transformer.h.45.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,314 - Quant - ERROR - Param tensor {transformer.h.45.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,322 - Quant - ERROR - Param tensor {transformer.h.46.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,335 - Quant - ERROR - Param tensor {transformer.h.46.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,345 - Quant - ERROR - Param tensor {transformer.h.47.ln_1.weight} not found in valid param set 2024-05-29 12:51:31,358 - Quant - ERROR - Param tensor {transformer.h.47.ln_2.weight} not found in valid param set 2024-05-29 12:51:31,366 - Quant - ERROR - Param tensor {transformer.ln_f.weight} not found in valid param set 2024-05-29 12:51:31,609 - Quant - INFO - Layers excluded from quantization: [] 2024-05-29 12:51:31,709 - AutoQuant - INFO - The results of Batchnorm Folding is saved in ./.trace/batchnorm_folding.pth and ./.trace/batchnorm_folding.encodings. 2024-05-29 12:51:31,735 - AutoQuant - INFO - Session finished: Batchnorm Folding. (eval score: -0.000019)`

Any suggestions on how to fix those warnings and errors ??