def _run_for_config(
        self, model: HfModelHandler, config: type[BasePassConfig], output_model_path: str
    ) -> HfModelHandler:
        """Run the selective mixed precision pass."""
        if not isinstance(model, HfModelHandler):
            raise ValueError("SelectiveMixedPrecision pass currently only supports HfModelHandler.")

        model_wrapper = ModelWrapper.from_model(model.load_model(cache_model=False))

        # TODO(jambayk): make this configurable
        override_config = {"bits": PrecisionBits.BITS8}
        overrides = {model_wrapper.get_lm_head()[1]: override_config}

        if config.algorithm == SelectiveMixedPrecision.Algorithm.K_QUANT_MIXED:
            layer_prefix = model_wrapper.get_layers()[1]
            num_layers = model_wrapper.num_hidden_layers
            for layer_idx, layer_wrapper in enumerate(model_wrapper.get_layer_wrappers()):
                if not (
                    layer_idx < num_layers / 8
                    or layer_idx >= 7 * num_layers / 8
                    or ((layer_idx - num_layers // 8) % 3 == 2)
                ):
                    continue

                # Add v_proj if exists
                attn_input_names = layer_wrapper.get_attention_inputs(return_name=True)[1]
                if len(attn_input_names) == 3:
                    overrides[f"{layer_prefix}.{layer_idx}.{attn_input_names[2]}"] = override_config
                # TODO(jambayk): how to handle models like phi-3/4 which have qkv merged?
                # what if downstream passes merge qkv for non-phi models?

                # Add down_proj
                mlp_output_names = layer_wrapper.get_mlp_outputs(return_name=True)[1]
                if len(mlp_output_names) == 1:
                    overrides[f"{layer_prefix}.{layer_idx}.{mlp_output_names[0]}"] = override_config

        # Create output model with mixed precision info as model attribute
        output_model = deepcopy(model)
        output_model.model_attributes = output_model.model_attributes or {}
        output_model.model_attributes["mixed_precision_info"] = {
            "default": {"bits": PrecisionBits.BITS4},
            "overrides": overrides,
        }
        return output_model
Run the selective mixed precision pass.
olive/passes/pytorch/selective_mixed_precision.py