From 6d717e63b8a2d3cbffe33d8aabd26aee3b0d1da4 Mon Sep 17 00:00:00 2001 From: Andrei Anufriev Date: Wed, 3 Jul 2024 16:33:24 +0200 Subject: [PATCH] Enabled scale estimation in compression. --- llm_bench/python/convert.py | 5 +++++ llm_bench/python/utils/conversion_utils/helpers.py | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/llm_bench/python/convert.py b/llm_bench/python/convert.py index eed79c794..00a9a8e62 100644 --- 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 index d7545950d..5b1f7bcd6 100644 --- 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"