puzzletron_nas_plugin

Puzzletron NAS plugin for the Modelopt framework (based on Puzzle algorithm: https://arxiv.org/abs/2411.19146).

It is used by mtn.convert() to convert a model from HF format to Puzzletron heterogeneous format + do pruning scoring and save pruned checkpoints, and by mtn.search() to perform the MIP-based NAS search.

Classes

PuzzletronModel

PuzzletronDescriptor

Descriptor for the Puzzletron mode.

PuzzletronSearcher

Runs NAS search for the Puzzletron mode.

Functions

convert_puzzletron_model

restore_puzzletron_model

Restore is not needed for the puzzletron mode as we are not saving any model state
ModeloptConfig PuzzletronConfig

Bases: ModeloptBaseConfig

Configuration for Puzzletron NAS algorithm.

Show default config as JSON
Default config (JSON):

{
   "input_model_path": "",
   "hydra_config_dir": "",
   "hydra_config_name": "",
   "puzzle_dir": "",
   "dataset_path": ""
}

field dataset_path: str
field hydra_config_dir: str
field hydra_config_name: str
field input_model_path: str
field puzzle_dir: str
class PuzzletronDescriptor

Bases: ModeDescriptor

Descriptor for the Puzzletron mode.

property config_class: type[ModeloptBaseConfig]

Configuration class for this mode.

property convert: Callable[[Module, ModeloptBaseConfig], tuple[Module, dict[str, Any]]] | Callable[[Module, ModeloptBaseConfig, Any], tuple[Module, dict[str, Any]]]

Entrypoint to convert a model.

property export_mode: str | None

The mode that corresponds to the export mode. For now, this will be a no-op as there is no modelopt’s concept of search space defined for the puzzletron algorithm.

property name: str

String identifier for this mode.

property restore: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], Module]

Entrypoint to restore a model.

property search_algorithm: type[BaseSearcher]

Return the associated searcher implementation.

class PuzzletronModel

Bases: Module

class PuzzletronSearcher

Bases: BaseSearcher

Runs NAS search for the Puzzletron mode.

property default_state_dict: dict[str, Any]

Not needed for the puzzletron mode as we are not saving any model state

Return type:

None

convert_puzzletron_model(model, config)

  1. Convert the model from HF format to AnyModel format.

  2. Score the pruning activations.

  3. Prune the model and save pruned checkpoints

The output of this step will be used by mnt.search() to perform the NAS search.

Parameters:
Return type:

tuple[Module, dict[str, Any]]

restore_puzzletron_model(model, config, metadata)

Restore is not needed for the puzzletron mode as we are not saving any model state

Parameters:
Return type:

Module