mode — Model Optimizer 0.29.0 (original) (raw)

Interface and utilities for optimization modes/algorithms.

A mode is a specific type or algorithms for model optimization, e.g., some type of algorithm for pruning or quantization. It can also specify a single step within an optimization algorithm instead of the whole algorithm. For example, a mode can prepare a model for pruning or export (i.e. fix the optimal model configuration) after pruning.

Within modelopt, a mode constitutes the unit for model optimization. We can define arbitrary modes, each mode gets recorded in the model’s modelopt state dict, and we can define workflows as a sequence of modes.

Classes

ModeDescriptor	Abstract class to describe a mode.

class ModeDescriptor

Bases: ABC

Abstract class to describe a mode.

assert_compatibility_as_next_mode_of(other_mode)

Assert that this mode is compatible as a next mode of the other mode.

Parameters:

other_mode (ModeDescriptor | str) –

Return type:

None

abstract property config_class_: type[ModeloptBaseConfig]_

Specifies the config class for the mode.

abstract property convert_: Callable[[Module, ModeloptBaseConfig], tuple[Module, dict[str, Any]]] | Callable[[Module, ModeloptBaseConfig, Any], tuple[Module, dict[str, Any]]]_

The mode’s entrypoint for converting a model.

The function signature of the convert entrypoint is described below:

Parameters:

model – Model to be restored.
config – Config used for the model that was also used during convert.
mode_kwargs – Additional keyword arguments to pass to the convert entrypoint.

Returns:

A tuple consisting of

the in-place modified model. If the modification failed, the entrypoint can return None instead
The config dict that can be used to call the restore entrypoint to instantly _restore_the modified model.
The metatdata that can be used to call the restore entrypoint to instantly_restore_ the modified model from the provided initial state, see below’s description for the restore entrypoint to get more info about metadata.

Raises:

ApplyModeError<modelopt.torch.opt._conversion.ApplyModeError>` to indicate that th –
conversion process failed. This error can be caught by user-facing APIs if they want to –
enable a fall-back behavior. –

property export_mode_: str | None_

The mode that corresponds to the export mode of this mode.

Certain modes require a subsequent export step. For example, after pruning, we might want to fine-tune the model and then export the model. This property specifies that mode if it exists.

None indicates that there exists no such mode.

Returns:

The (optional) mode name that corresponds to the export mode of this mode. Defaults to None.

property is_export_mode_: bool_

Whether the mode is an export mode.

Returns:

True if the mode is an export mode, False otherwise. Defaults to False.

abstract property name_: str_

Returns the string name of the mode.

property next_modes_: set[str] | None_

Modes that must immediately follow this mode.

Certain modes only makes sense if they are followed by certain other modes.

An empty set indicates that _no_ mode can follow this mode. A None value indicates that there are no restrictions on the following mode.

Returns:

A set of mode names that must immediately follow this mode. Defaults to None.

property next_prohibited_modes_: set[str] | None_

Modes that should not be applied after this mode.

This is used to specify that certain modes should not be applied after this mode. A None value indicates that incompatibitly check is skipped for this mode. This is useful if there are many allowed modes but only a few prohibited modes.

Returns:

A set of mode names that should not be applied after this mode. Defaults to None.

property require_model_like_: bool_

Whether the mode requires a ModelLike input?

Returns:

False

abstract property restore_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], Module]_

The mode’s entrypoint for restoring a model.

The function signature of the restore entrypoint is described below:

Parameters:

model – Model to be restored.
config – Config used for the model that was also used during convert.
metadata – The metadata is used during restoration of the model architecture to instantly restore the modified model. The metadata is used on top of the config to ensure that the model can be instantly restored/modified from the provided state. This is helpful when the convert entrypoint contains non-deterministic operations whose outcome can be stored in the metadata to ensure that the model can be restored reliably.A few examples of potential non-deterministic operations are provided below: (1) Latency measurements: if the conversion leverages latency measurements during conversion the conversion process may become non-deterministic. (2) Random operations: if the conversion leverages random operations during conversion, we should store the samples or random seed. (3) Module’s train flag: the conversion process might be affected by the module’s train flag (e.g. tracing is indirectly affected by train flag since the forward may be affected by the train flag). If so, we should store the train flag in the metadata and set the model into the correct mode.

Returns:

The in-place modified and restored model. If the modification failed, the entrypoint can

return None instead

Raises:

ApplyModeError<modelopt.torch.opt._conversion.ApplyModeError>` to indicate that th –
conversion process failed. This error can be caught by user-facing APIs if they want to –
enable a fall-back behavior. –

property search_algorithm_: type[BaseSearcher]_

Specifies the search algorithm to use for this mode (if any).

property update_for_new_mode_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]_

The mode’s (optional) entrypoint for updating a model’s config and metadata before a new mode.

This is useful if metadata or config needs to be updated before adding a new mode. For example, a after adding a new mode, the current mode’s restore might only need a subset of the metadata/config.

The function signature of this update entrypoint is described below:

Parameters:

model – Model to be restored.
config – The config as described above. It should be modified IN-PLACE.
metadata – The metadata as described above. It should be modified IN-PLACE.

Returns:

None.

property update_for_save_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]_

The mode’s (optional) entrypoint for updating a model’s config and metadata before saving.

This is useful if metadata or config needs to be updated for saving (and restoring) the mode.

The function signature of this update entrypoint is described below:

Parameters:

model – Model to be restored.
config – The config as described above. It should be modified IN-PLACE.
metadata – The metadata as described above. It should be modified IN-PLACE.

Returns:

None.