Implement pass search #1562

shaahji · 2025-01-21T19:35:08Z

Implement pass search

Reimplement search logic to include passes in search space.

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.
Is this PR including examples changes? If yes, please remember to update example documentation in a follow-up PR.

(Optional) Issue link

review-notebook-app · 2025-01-21T19:35:14Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

test/unit_test/search/test_search_strategy.py

examples/mobilenet/prepare_config.py

olive/engine/engine.py

olive/search/samplers/__init__.py

olive/search/search_point.py

olive/search/search_strategy.py

olive/workflows/run/config.py

test/unit_test/search/test_search_results.py

test/unit_test/search/test_search_strategy.py

olive/search/search_parameter.py

olive/search/search_results.py

olive/search/search_strategy.py

examples/test/local/test_bert_cuda_gpu.py

xiaoyu-work · 2025-01-23T00:09:48Z

Can you also update documents related to your changes?

examples/bert/user_script.py

@@ -8,7 +8,7 @@
 import torch
 import torchmetrics
 import transformers
-from datasets import load_dataset, load_metric
+from datasets import load_dataset


examples/bert/user_script.py

+try:
+    from datasets import load_metric
+except ImportError:
+    from evaluate import load as load_metric


test/unit_test/search/test_search_strategy.py

+import re
+from collections import OrderedDict
+
+import pytest


Reimplement search logic to include passes in search space.

shaahji · 2025-01-23T10:36:25Z

Can you also update documents related to your changes?

Done!

jambayk · 2025-01-23T22:44:20Z

examples/__init__.py

why is this required?

jambayk · 2025-01-23T22:46:18Z

examples/bert/bert.py

please update the "BERT optimization with CUDA/TensorRT on GPU" section in the readme

jambayk · 2025-01-23T22:55:56Z

examples/llama2/llama2.py

    if not args.use_gptq:
-        template_json["pass_flows"] = [flow for flow in SUPPORTED_WORKFLOWS[device] if "gptq" not in flow[0]]
+        used_passes = [


I think we might have to create a new mapping from gptq/no-gptq , precision, etc.
Previously this resulted in multiple pass flows. But now, you are just flattening the pass flows into a single list.

I think the passes used might looks like "conversion_merged", "transformers_optimization_fp16", "conversion_merged", "transformers_optimization_fp16", "blockwise_quant_int4" which is not the intended behavior here. You can test this by using the --config_only option to dump the config.

we need to either make the mapping specific or generate multiple workflows and run the separately.

jambayk · 2025-01-23T23:04:55Z

examples/phi2/phi2.py

@@ -186,25 +186,23 @@ def main(raw_args=None):
            legacy_optimization_setting(template_json)

        # add pass flows
-        pass_flows = [[]]
+        used_passes = {}


Just in general, I think using sets might be a bit risky since they are unordered.

jambayk · 2025-01-23T23:07:07Z

examples/phi3/phi3.py

@@ -191,7 +191,10 @@ def use_passes(template_json, *passes):
    else:
        del template_json["data_configs"]

-    template_json["pass_flows"] = [passes]
+    for pass_name in set(template_json["passes"].keys()):


I think instead of popping the unused passes, it might be better to create a new dict with the used passes. Popping assumes the order in template_json["passes"] is the same as in passes

jambayk · 2025-01-23T23:10:22Z

examples/stable_diffusion/sd_utils/ort.py

@@ -23,7 +23,10 @@ def update_cuda_config(config_cuda: Dict):
    if version.parse(OrtVersion) < version.parse("1.17.0"):
        # disable skip_group_norm fusion since there is a shape inference bug which leads to invalid models
        config_cuda["passes"]["optimize_cuda"]["optimization_options"] = {"enable_skip_group_norm": False}
-    config_cuda["pass_flows"] = [["convert", "optimize_cuda"]]
+    used_passes = {"convert", "optimize_cuda"}


this part in stable_diffusion.py also needs to be updated to only use "convert", "optimize"

Olive/examples/stable_diffusion/stable_diffusion.py

Line 181 in 28beb3f

if provider == "dml":

jambayk · 2025-01-23T23:23:11Z

olive/search/samplers/optuna_sampler.py

+
+        # Initialize the searcher
+        self._sampler = self._create_sampler()
+        # TODO(olivedev): There is no absolute direction to set.


will this be investigated later? Since I still think we need directions for the signals to make sense and the sampler to choose the next option.

jambayk · 2025-01-23T23:25:19Z

olive/search/samplers/optuna_sampler.py

+        if self.should_stop:
+            return None
+
+        while True:


is this to avoid sampling the same point? If so, could you add a comment about it?
does it do the same even with the current implementation of optuna sampler?

jambayk · 2025-01-23T23:33:08Z

olive/search/search_strategy.py

+
+
+@dataclass
+class SearchWalkState:


could you add some docstrings to describe what each class does/is used for? thanks!

jambayk · 2025-01-23T23:34:24Z

olive/search/search_strategy.py

+        self._init_model_id: str = None
+
+        # State variables
+        self._path: List[int] = None


SearchWalkState also has a path attribute

jambayk · 2025-01-23T23:51:26Z

olive/search/samplers/optuna_sampler.py

+            suggestion_index = trial.suggest_categorical(suggestion_name, list(range(suggestion_len)))
+            suggestion = suggestions[suggestion_index]
+
+            if isinstance(suggestion, (SearchParameter, SearchSpace)):


does this case happen? I thought this the if block at line 87 is the base case. so it should only result in a fixed value?

jambayk · 2025-01-23T23:55:09Z

olive/search/samplers/optuna_sampler.py

+
+            spi = 0
+            for child_index, suggestions_len in reversed(indicies_lengths):
+                spi *= suggestions_len


does this guarantee uniqueness for the search point index? behaves like a generic version of binary (some base_n) encoding?

olive/search/samplers/search_sampler.py

jambayk · 2025-01-24T00:11:56Z

olive/search/search_space.py

+            index, values[name] = SearchSpace.get_suggestion(param, index, values)
+        return values
+
+    @staticmethod


could you add some docstrings and comments to describe the functionalities and logic. It's a bit hard to follow

shaahji mentioned this pull request Jan 21, 2025

Validate pass config before instantiating the pass #1553

Merged

6 tasks

github-advanced-security bot found potential problems Jan 21, 2025

View reviewed changes

test/unit_test/search/test_search_strategy.py Fixed Show resolved Hide resolved

github-advanced-security bot found potential problems Jan 21, 2025

View reviewed changes

shaahji force-pushed the shaahji/pass_search branch 3 times, most recently from 64d4138 to 534a692 Compare January 21, 2025 22:47

xiaoyu-work reviewed Jan 22, 2025

View reviewed changes

shaahji force-pushed the shaahji/pass_search branch 2 times, most recently from 3e6fbfd to 457d3b7 Compare January 23, 2025 09:52

github-advanced-security bot found potential problems Jan 23, 2025

View reviewed changes

shaahji force-pushed the shaahji/pass_search branch 2 times, most recently from a87fd30 to 58fa43e Compare January 23, 2025 10:27

Implement pass search

d63c35a

Reimplement search logic to include passes in search space.

shaahji force-pushed the shaahji/pass_search branch from 58fa43e to d63c35a Compare January 23, 2025 10:31

jambayk reviewed Jan 23, 2025

View reviewed changes

examples/__init__.py

Copy link

Contributor

jambayk Jan 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this required?

jambayk reviewed Jan 23, 2025

View reviewed changes

examples/bert/bert.py

Copy link

Contributor

jambayk Jan 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please update the "BERT optimization with CUDA/TensorRT on GPU" section in the readme

jambayk reviewed Jan 23, 2025

View reviewed changes

olive/search/samplers/search_sampler.py Show resolved Hide resolved

jambayk reviewed Jan 24, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement pass search #1562

Implement pass search #1562

shaahji commented Jan 21, 2025 •

edited

Loading

review-notebook-app bot commented Jan 21, 2025

xiaoyu-work commented Jan 23, 2025

shaahji commented Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025 •

edited

Loading

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025 •

edited

Loading

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 23, 2025

jambayk Jan 24, 2025

Implement pass search #1562

Are you sure you want to change the base?

Implement pass search #1562

Conversation

shaahji commented Jan 21, 2025 • edited Loading