Implement pass search (#1562)

## Implement pass search Reimplement search logic to include passes in search space. ## Checklist before requesting a review - [x] Add unit tests for this change. - [x] Make sure all tests can pass. - [ ] Update documents if necessary. - [x] 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](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md) in a follow-up PR. ## (Optional) Issue link
microsoft · Jan 31, 2025 · c22f68b · c22f68b
1 parent e314b00
commit c22f68b
Show file tree

Hide file tree

Showing 126 changed files with 5,177 additions and 1,525 deletions.
diff --git a/.gitignore b/.gitignore
@@ -126,6 +126,7 @@ celerybeat.pid
 # Environments
 .env
 .venv
+.vs
 env/
 venv/
 ENV/

diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -117,37 +117,6 @@ init_command = [
   '--requirement=requirements-dev.txt'
 ]
 
-[[linter]]
-code = 'MYPY'
-include_patterns = [
-  '**/*.py',
-  '**/*.pyi'
-]
-exclude_patterns = [
-  'examples/pytorch/*.py'
-]
-command = [
-  'python',
-  '-m',
-  'lintrunner_adapters',
-  'run',
-  'mypy_linter',
-  '--config=pyproject.toml',
-  '--show-notes',
-  '--show-disable',
-  '--',
-  '@{{PATHSFILE}}'
-]
-init_command = [
-  'python',
-  '-m',
-  'lintrunner_adapters',
-  'run',
-  'pip_init',
-  '--dry-run={{DRYRUN}}',
-  'mypy==1.0.0'
-]
-
 [[linter]]
 code = 'NOQA'
 include_patterns = ['**/*.py', '**/*.pyi']

diff --git a/docs/architecture.md b/docs/architecture.md
@@ -13,12 +13,12 @@ This document describes the Olive components, and some implementation details. T
 - [Search](#search)
     - [Search Parameter](#searchparameter)
     - [Search Space](#searchspace)
-    - [Search Algorithm](#searchalgorithm)
+    - [Search Sampler](#searchsampler)
     - [Search Results](#searchresults)
-- [Search Strategy](#search-strategy)
-    - [Execution order](#execution-order)
-    - [User Interface](#user-interface)
-    - [Implementation](#implementation)
+    - [Search Strategy](#search-strategy)
+        - [Execution order](#execution-order)
+        - [User Interface](#user-interface)
+        - [Implementation](#implementation)
 - [System](#system)
     - [OliveSystem Class](#olivesystem-class)
 - [Data Container](#data-container)
@@ -109,9 +109,9 @@ The engine maintains a cache directory with three sub-directories:
 - `mlflow`: stores mlflow model files.
 
 ## Search
-Olive workflows support search parameters which are optimized using search algorithms.
+Olive workflows support search parameters which are optimized using various execution order.
 
-At the most basic level is `SearchParameter` which describes the options for search parameters. `SearchSpace` combines search parameters for one or more passes and `SearchAlgorithm` provides different sampling algorithms to search for the best parameter configuration (search point) from the search space.
+At the most basic level is `SearchParameter` which describes the options for search parameters. `SearchSpace` combines search parameters for one or more passes and `SearchSampler` provides different sampling algorithms to search for the best parameter configuration (search point) from the search space.
 
 ### SearchParameter
 A search parameter defines a discrete categorical distribution.
@@ -122,7 +122,7 @@ There are two types of search parameters:
 
 **Note:**
 - There cannot be any cyclic parent child dependencies.
-- Search algorithms order the search parameters topologically so that the parents are sampled before the children.
+- Search space orders the search parameters topologically so that the parents are sampled before the children.
 
 ### SearchSpace
 Search space combines search parameters from one or more passes and provides methods to iterate over the search space (`iterate`) or generate random samples (`random_sample`).
@@ -139,18 +139,18 @@ The corresponding conceptual search space is the space of all possible parameter
 {“pass_id/space_name”: {“param_name”: param_value}}
 ```
 
-### SearchAlgorithm
-Search algorithm operates over a search space and provides samples/trials (search points) from the search space to execute and evaluate.
+### SearchSampler
+Sampling algorithm queries over a search space and provides samples/trials (search points) from the search space to evaluate.
 
-Each search algorithm provides the methods:
+Each search sampler provides the methods:
 - `suggest`: returns a search point to execute and evaluate. The algorithm can sample a search point based on the evaluation results for previously suggested points.
 - `report`: report evaluation results for a search point. The search point can also be pruned if it contains invalid pass configs or failed during execution/evaluation.
 
-The following search algorithms have been implemented:
-- `ExhaustiveSearchAlgorithm`: Exhaustively iterates over the search space.
-- `RandomSearchAlgorithm`: Randomly samples points from the search space without replacement.
-- `OptunaSearchAlgorithm`: Abstract base class for algorithms built using `optuna` samplers. This class cannot be used directly
-    - `TPESearchAlgorithm`: Uses optuna `TPESampler`.
+The following sampling algorithms have been implemented:
+- `SequentialSampler`: Sequentially iterates over the search space.
+- `RandomSampler`: Randomly samples points from the search space.
+- `OptunaSampler`: Abstract base class for algorithms built using `optuna` samplers.
+    - `TPESampler`: Uses optuna's `TPESampler`.
 
 ### SearchResults
 `SearchResults` stores evaluation results for samples made from a search space and provides tools to analyze and select the best search point/s.
@@ -159,10 +159,10 @@ Results are reported using the `record` method.
 
 Currently `best_search_point` selects the best search point by maximizing/minimizing metrics using tie breaking. We intend to provide different model selection strategies for both single and multi-objective optimization.
 
-## Search Strategy
+## SearchStrategy
 Search strategy provides an optimization pipeline that finds the best search point from the search space of one or more passes.
 
-It consists of two sub-components – `execution_order` and `search_algorithm`. Search algorithm has been covered in the previous section.
+It consists of two sub-components – `execution_order` and `sampler`. Sampling algorithms have been covered in the previous section.
 
 ### Execution Order
 The execution order defines the order in which the passes are optimized.

diff --git a/docs/source/extending/design.md b/docs/source/extending/design.md
@@ -5,8 +5,7 @@ that are composed to construct a model optimization workflow.
 The workflow which is run by the **Engine** is composed of **Passes** that are executed in a specific order.
 Each Pass is responsible for performing a specific optimization on the model. Each Pass might have a set of parameters that
 can be tuned to achieve the best metrics, say accuracy and latency, that are evaluated by the respective **Evaluator**.
-The Engine employs a **Search Strategy** that uses a **Search Algorithm** to auto-tune each Pass one by one or set of Passes
-together.
+The Engine employs a **Search Strategy** that uses a **Search Samplers** to auto-tune each Pass one by one or set of Passes together.
 
 Each Pass can be run on any host **System** and its output model can be evaluated on the desired target **System**.
 
@@ -71,22 +70,23 @@ created and **registered** along with their host system and evaluators if any.
 The engine also maintains a cache directory to cache pass runs, models and evaluations.
 
 ## Search Strategy
-Search strategy provides an optimization pipeline that finds the best search point from the search space of one or more passes.
+Search strategy provides an optimization pipeline that finds the best search point from the search space built from one or more passes and search parameters within each of those passes. `include_pass_params` controls whether or not individual pass' search parameters are included in the search. `max_iter` and `max_time` can be configured for finer control.
 
-It consists of two sub-components – `execution_order` and `search_algorithm`.
+It consists of two sub-components – `execution_order` and `sampler`.
 
 ### Execution Order
-The execution order defines the order in which the passes are optimized.
+The execution order defines the order in which the search space is traversed.
 
 Currently, we support two execution orders:
-- `joint`: The search spaces of all passes are combined and searched together to find the best search point. Each search point
-that is evaluated has parameters for the search parameters of all passes.
-- `pass-by-pass`: The search space of each pass is searched and optimized independently in order.
+- `joint`: The search spaces of all passes and their corresponding parameters are combined and searched together to find the best search point. Each search point consists of values for all search parameters of at most one pass in each pass group.
+- `pass-by-pass`: The search space of each pass group is searched and optimized independently in order.
 
-### Search Algorithm
-Search algorithm operates over a search space and provides samples/trials (search points) from the search space to execute and evaluate.
+### Search Sampler
+Search sampler provides samples/trials (search points) from the search space to evaluate. Each search point consists of valuesfor all search parameters and all passes within the pass group.
 
-The following search algorithms have been implemented:
-- `exhaustive`: Exhaustively iterates over the search space.
-- `random`: Randomly samples points from the search space without replacement.
-- `tpe`: ample using TPE (Tree-structured Parzen Estimator) algorithm to sample from the search space.
+The following sampling algorithms have been implemented:
+- `sequential`: Sequentially iterates over the search space.
+- `random`: Randomly samples points from the search space.
+- `tpe`: Sample using TPE (Tree-structured Parzen Estimator) algorithm to sample the search space.
+
+Each of the sampler can be used as an exhaustive search by setting the `max_samples` field to zero.
diff --git a/docs/source/how-to/configure-workflows/auto-opt.md b/docs/source/how-to/configure-workflows/auto-opt.md
@@ -41,11 +41,9 @@ Here is a simple example of Auto Optimizer configuration, the item which is not
         "engine": {
             "search_strategy": {
                 "execution_order": "joint",
-                "search_algorithm": "tpe",
-                "search_algorithm_config": {
-                    "num_samples": 1,
-                    "seed": 0
-                }
+                "sampler": "tpe",
+                "max_samples": 1,
+                "seed": 0
             },
             "evaluator": "common_evaluator",
             "cache_dir": "cache",
@@ -149,8 +147,8 @@ Here is another quick comparison between Auto Optimizer and manual settings.
     },
     "search_strategy": {
         "execution_order": "joint",
-        "search_algorithm": "tpe",
-        "num_samples": 1,
+        "sampler": "tpe",
+        "max_samples": 1,
         "seed": 0
     },
     "evaluator": "common_evaluator",
@@ -261,8 +259,8 @@ Here is another quick comparison between Auto Optimizer and manual settings.
     ],
     "search_strategy": {
         "execution_order": "joint",
-        "search_algorithm": "tpe",
-        "num_samples": 1,
+        "sampler": "tpe",
+        "max_samples": 1,
         "seed": 0
     },
     "evaluator": "common_evaluator",

diff --git a/docs/source/how-to/configure-workflows/model-packaging.md b/docs/source/how-to/configure-workflows/model-packaging.md
@@ -197,11 +197,9 @@ You can add different types `PackagingConfig` as a list to Engine configurations
 "engine": {
     "search_strategy": {
         "execution_order": "joint",
-        "search_algorithm": "tpe",
-        "search_algorithm_config": {
-            "num_samples": 5,
-            "seed": 0
-        }
+        "sampler": "tpe",
+        "max_samples": 5,
+        "seed": 0
     },
     "evaluator": "common_evaluator",
     "host": "local_system",

diff --git a/docs/source/reference/index.rst b/docs/source/reference/index.rst
@@ -33,11 +33,11 @@ Reference
       :octicon:`arrow-right;1em;sd-text-info` `Pass <pass.html>`_
 
    .. grid-item-card::
-      **Search Algorithm**
+      **Search Samplers**
 
-      Configure search strategies.
+      Configure search samplers.
 
-      :octicon:`arrow-right;1em;sd-text-info` `Search Algorithm <search-algorithm.html>`_
+      :octicon:`arrow-right;1em;sd-text-info` `Search Samplers <search-samplers.html>`_
 
 
 .. toctree::
@@ -48,4 +48,4 @@ Reference
    options
    model
    pass
-   search-algorithm
+   search-samplers
diff --git a/docs/source/reference/options.md b/docs/source/reference/options.md
@@ -446,41 +446,6 @@ Please also find the detailed options from following table for each pass:
 }
 ```
 
-## Pass Flows Information
-
-`pass_flows: List[List[str]]`
-
-This is a list of list of pass names. Each list of pass names is a pass flow which will be executed in order.
-When `pass_flows` is not specified, the passes are executed in the order of the `passes` dictionary.
-
-### Example
-
-```json
-"passes": {
-    "onnx_conversion": {
-        "type": "OnnxConversion",
-        "target_opset": 13
-    },
-    "transformers_optimization": {
-        "type": "OrtTransformersOptimization",
-        "model_type": "bert",
-        "num_heads": 12,
-        "hidden_size": 768,
-        "float16": true
-    },
-    "onnx_quantization": {
-        "type": "OnnxQuantization",
-        "data_config": "calib_data_coonfig",
-        "weight_type": "QUInt8"
-    }
-},
-"pass_flows": [
-    ["onnx_conversion", "transformers_optimization"],
-    ["onnx_conversion", "transformers_optimization", "onnx_quantization"],
-    ["onnx_conversion", "onnx_quantization"],
-]
-```
-
 ## Engine Information
 
 `engine: [Dict]`
@@ -491,15 +456,13 @@ This is a dictionary that contains the information of the engine. Its fields can
 
     - `execution_order: [str]` The execution order of the optimizations of passes. The options are `pass-by-pass` and `joint`.
 
-    - `search_algorithm: [str]` The search algorithm of the engine. The available search algorithms are `exhaustive`, `random` and `tpe`.
+    - `sampler: [str]` The search sampler to use while traversing. The available search algorithms are `random`, `sequential` and `tpe`.
 
-    - `search_algorithm_config: [Dict]` The configuration of the search algorithm. The configuration of the search algorithm depends on
-    the search algorithm. Its fields can be provided directly to the parent dictionary.
+    - `sampler_config: [Dict]` The configuration of the sampler. The options depends on the chosen sampler. Its fields can be provided directly to the parent dictionary.
 
-    - `output_model_num: [int]` The number of output models from the engine based on metric priority. If not specified, the engine will output all qualified models.
+    - `stop_when_goals_met: [Boolean]` This decides whether to stop the search when the metric goals, if any,  are met. This is `false` by default.
 
-    - `stop_when_goals_met: [Boolean]` This decides whether to stop the search when the metric goals, if any,  are met. This is `false` by
-    default.
+    - `include_pass_params: [Boolean]` Includes individual pass parameter to build the search space. Defaults to true.
 
     - `max_iter: [int]` The maximum number of iterations of the search. Only valid for `joint` execution order. By default, there is no
     maximum number of iterations.
@@ -511,8 +474,8 @@ This is a dictionary that contains the information of the engine. Its fields can
   have empty search spaces. The output of the final pass will be evaluated if there is a valid evaluator. The output of the engine will be
   the output model of the final pass and its evaluation result.
 
-  If `search_strategy` is `true`, the search strategy will be the default search strategy. The default search strategy is `exhaustive` search
-  algorithm with `joint` execution order.
+  If `search_strategy` is `true`, the search strategy will be the default search strategy. The default search strategy is `sequential` search
+  sampler with `joint` execution order.
 
 - `evaluate_input_model: [Boolean]` In this mode, the engine will evaluate the input model using the engine's evaluator and return the results. If the engine has no evaluator, it will skip the evaluation. This is `true` by default.
 
@@ -556,22 +519,24 @@ This is a dictionary that contains the information of the engine. Its fields can
 - `log_to_file: [Boolean]`, `false` by default. This decides whether to log to file. If `true`, the log will be stored in a olive-<timestamp>.log file
     under the current working directory.
 
-Please find the detailed config options from following table for each search algorithm:
+Please find the detailed config options from following table for each search sampler:
+
+Note that if `max_samples` is set to zero, each of the below sampler will be exhaustive.
 
-| Algorithm  | Description |
+| Sampler  | Description |
 |:----------|:-------------|
-| [exhaustive](../../reference/search-algorithm.rst#_exhaustive_search_algorithm) | Iterates over the entire search space |
-| [random](../../reference/search-algorithm.rst#_random_search_algorithm) | Samples random points from the search space with or without replacement |
-| [tpe](../../reference/search-algorithm.rst#_tpe_search_algorithm) | Sample using TPE (Tree-structured Parzen Estimator) algorithm. |
+| [random](../../reference/search-samplers.rst#_random_sampler) | Samples random points from the search space |
+| [sequential](../../reference/search-samplers.rst#_sequential_sampler) | Iterates over the entire search space sequentially |
+| [tpe](../../reference/search-samplers.rst#_tpe_sampler) | Sample using TPE (Tree-structured Parzen Estimator) algorithm. |
 
 ### Example
 
 ```json
 "engine": {
     "search_strategy": {
         "execution_order": "joint",
-        "search_algorithm": "tpe",
-        "num_samples": 5,
+        "sampler": "tpe",
+        "max_samples": 5,
         "seed": 0
     },
     "evaluator": "common_evaluator",

diff --git a/docs/source/reference/search-algorithm.rst b/docs/source/reference/search-algorithm.rst
-Original file line number
+Diff line change
@@ Expand Up / @@ -126,6 +126,7 @@ celerybeat.pid @@
     # Environments
     .env
     .venv
+    .vs
     env/
     venv/
     ENV/
@@ Expand Down @@