diff --git a/docs/source/reference/python_api.md b/docs/source/reference/python_api.md index 8bf3d40a0f..414755d774 100644 --- a/docs/source/reference/python_api.md +++ b/docs/source/reference/python_api.md @@ -39,7 +39,7 @@ Optimize the input model with comprehensive pass scheduling. - `act_precision` (str, optional): Activation precision for quantization. - `num_split` (int, optional): Number of splits for model splitting. - `memory` (int, optional): Available device memory in MB. -- `exporter` (str, optional): Exporter to use ("model_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"). +- `exporter` (str, optional): Exporter to use ("model_builder", "mobius_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"). - `dim_param` (str, optional): Dynamic parameter names for dynamic to fixed shape conversion. - `dim_value` (str, optional): Fixed dimension values for dynamic to fixed shape conversion. - `use_qdq_format` (bool): Use QDQ format for quantization. Defaults to `False`. diff --git a/olive/cli/optimize.py b/olive/cli/optimize.py index 06aa92ca06..f25380335a 100644 --- a/olive/cli/optimize.py +++ b/olive/cli/optimize.py @@ -113,7 +113,7 @@ def register_subcommand(parser: ArgumentParser): sub_parser.add_argument( "--exporter", type=str, - choices=["model_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"], + choices=["model_builder", "mobius_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"], help="Exporter to use for model conversion (optional).", ) @@ -198,6 +198,7 @@ def __init__(self, parser: ArgumentParser, args: Namespace, unknown_args: Option self.enable_gptq = False self.enable_capture_split_info = False self.enable_model_builder = False + self.enable_mobius_builder = False self.enable_onnx_conversion = False self.enable_optimum_openvino_conversion = False self.enable_dynamic_to_fixed_shape = False @@ -256,6 +257,16 @@ def _validate_arguments(self): if self.args.exporter is None and self.args.modality == "text": self.args.exporter = "model_builder" + if self.args.exporter == "mobius_builder" and self.args.precision not in ( + Precision.FP32, + Precision.FP16, + Precision.BF16, + ): + raise ValueError( + f"MobiusBuilder supports precisions fp32/fp16/bf16; got '{self.args.precision}'. " + "For INT4, capture in fp32/fp16/bf16 first and run a quantization pass afterwards." + ) + if self.args.modality not in ["text"]: raise ValueError(f"Unsupported modality: {self.args.modality}. Only 'text' is supported for optimization.") @@ -336,6 +347,10 @@ def _build_passes_config(self) -> dict[str, Any]: if self.enable_model_builder: passes_config["model_builder"] = self._get_model_builder_pass_config() + self.enable_mobius_builder = self._enable_mobius_builder_pass() + if self.enable_mobius_builder: + passes_config["mobius_builder"] = self._get_mobius_builder_pass_config() + self.enable_onnx_conversion = self._enable_onnx_conversion_pass() if self.enable_onnx_conversion: passes_config["onnx_conversion"] = self._get_onnx_conversion_pass_config() @@ -502,6 +517,19 @@ def _get_model_builder_pass_config(self) -> dict[str, Any]: return config + def _enable_mobius_builder_pass(self) -> bool: + """Return true if condition to add MobiusBuilder pass is met.""" + provider = ExecutionProvider(self.args.provider) + return ( + self.is_hf_model + and provider != ExecutionProvider.OpenVINOExecutionProvider + and self.args.exporter == "mobius_builder" + ) + + def _get_mobius_builder_pass_config(self) -> dict[str, Any]: + """Return pass dictionary for MobiusBuilder pass.""" + return {"type": "MobiusBuilder", "precision": Precision(self.args.precision).value} + def _enable_onnx_conversion_pass(self) -> bool: """Return true if condition to add OnnxConversion pass is met.""" provider = ExecutionProvider(self.args.provider) @@ -564,7 +592,7 @@ def _get_openvino_io_update_pass_config(self) -> dict[str, Any]: def _enable_onnx_peephole_optimizer_pass(self) -> bool: """Return true if condition to add OnnxPeepholeOptimizer pass is met.""" - return not self.is_hf_model or self.args.exporter != "model_builder" + return not self.is_hf_model or self.args.exporter not in ["model_builder", "mobius_builder"] def _get_onnx_peephole_optimizer_pass_config(self) -> dict[str, Any]: """Return pass dictionary for OnnxPeepholeOptimizer pass.""" @@ -634,7 +662,7 @@ def _get_onnx_blockwise_rtn_quantization_pass_config(self) -> dict[str, Any]: def _enable_onnx_float_to_float16_pass(self) -> bool: """Return true if condition to add OnnxFloatToFloat16 pass is met.""" precision = Precision(self.args.precision) - return precision == Precision.FP16 and not self.enable_model_builder + return precision == Precision.FP16 and not (self.enable_model_builder or self.enable_mobius_builder) def _get_onnx_float_to_float16_pass_config(self) -> dict[str, Any]: """Return pass dictionary for OnnxFloatToFloat16 pass.""" diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 3db547f63d..2e74b441b7 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -840,6 +840,13 @@ def test_optimize_cli_pass_list(mock_repo_exists, mock_run, tmp_path): None, "CUDAExecutionProvider", ], + [ + "optimize", + "--precision fp16 --provider CUDAExecutionProvider --exporter mobius_builder", + "MobiusBuilder", + None, + "CUDAExecutionProvider", + ], [ "optimize", (