Enabled scale estimation in compression. (#575)

openvinotoolkit · Jul 3, 2024 · 119de13 · 119de13
1 parent 98481c7
commit 119de13
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/llm_bench/python/convert.py b/llm_bench/python/convert.py
@@ -1449,6 +1449,11 @@ def main():
         action="store_true",
         help="Apply AWQ algorithm during compression",
     )
+    compression_group.add_argument(
+        "--scale_estimation",
+        action="store_true",
+        help="Apply scale estimation algorithm during compression",
+    )
     add_stateful_model_arguments(parser)
 
     args = parser.parse_args()

diff --git a/llm_bench/python/utils/conversion_utils/helpers.py b/llm_bench/python/utils/conversion_utils/helpers.py
@@ -160,10 +160,14 @@ def get_data_aware_args(ov_model, tokenizer, config, compression_args, args):
             res['mode'] = dataset_args['sensitivity_metric']
         if 'awq' in dataset_args:
             res['awq'] = dataset_args['awq']
+        if 'scale_estimation' in dataset_args:
+            res['scale_estimation'] = dataset_args['scale_estimation']
     elif args.dataset is not None:
         dataset_params = args.dataset
         if args.awq:
             res['awq'] = args.awq
+        if args.scale_estimation:
+            res['scale_estimation'] = args.scale_estimation
 
     if dataset_params is not None:
         # for example "wikitext,wikitext-2-v1,train[:1000],text"