dmx.compressor.modeling.model

Classes

`DmxConfig`	This is a dict of Dmx-specific configurations for a DmxModel This defines the 'states' to be optimized
`DmxConfigRule`([module_types, name_re, ...])	This is a rule that specifies how to transform from DmxConfig to DmxConfig This defines the 'action' in the state space
`DmxModel`()
`DmxModelMixin`()
`DmxPipelineMixin`()
`DmxSimplePipeline`(model_dict[, preproc, ...])
`DmxTransformation`	alias of `DmxConfigRule`
`Model`(body[, head, tail, hf, concrete_args])	This is a backward-compatible placeholder for legacy models.

class dmx.compressor.modeling.model.DmxConfig

Bases: dict

This is a dict of Dmx-specific configurations for a DmxModel This defines the ‘states’ to be optimized

classmethod from_model(model: Module, freeze=False)

A function that stores state and ops format of the model in a DmxConfig object

Parameters:: model (Model) – Target model for creating the DmxConfig
Returns:: A DmxConfig object that stores state and ops format of the model in a DmxConfig object

classmethod from_yaml(fname)

A function that creates a DmxConfig object from a yaml file

Parameters:: fname (str) – file path of the yaml file
Returns:: A DmxConfig object vreated from yaml file

property module_names: Returns the module names in the DmxConfig object

to_yaml(fname)

A function that stores the content of a DmxConfig object to a yaml file

Parameters:: fname (str) – file path of the yaml file

class dmx.compressor.modeling.model.DmxConfigRule(module_types=(), name_re: str = '', module_config: DmxModuleConfig = {})

Bases: SimpleNamespace

This is a rule that specifies how to transform from DmxConfig to DmxConfig This defines the ‘action’ in the state space

Parameters:

module_types (Optional[Tuple]) – Types of modules to apply transformation. Defaults to empty Tuple.
name_re (Optional[str]) – String for matching module name patterns. Defaults to empty str.
module_config (Optional[DmxModuleConfig]) – DmxModuleConfig that specifies the ops formats for a module. Defaults to empty DmxModuleConfig.

module_types

Types of modules to apply transformation.

Type:: Tuple

name_rule

Pattern for matching module names.

Type:: Pattern

module_config

DmxModuleConfig that specifies the ops formats for a module.

Type:: DmxModuleConfig

apply_to(model_or_config: Module | DmxConfig)

A function that sets format of ops according to self.module_config for modules selected by self.module_types and self.name_rule on a model or DmxConfig

Parameters:: model_or_config (Union[Model, DmxConfig]) – Model or DmxConfig to apply transformation on.

names_in(model_or_config: Module | DmxConfig)

Creates a list of module names where the modules are in self.module_types and the names match with self.name_rule.

Parameters:: model_or_config (Union[torch.nn.Module, DmxConfig]) – torch.nn.Module or DmxConfig to create the name of modules for.
Returns:: A list of module names

class dmx.compressor.modeling.model.DmxModel

Bases: DmxModelMixin

classmethod from_torch(model: Module, additional_dmx_aware_mappings=None, export=False) → Module

make_compiler_graphs()

visualize_graph(out_file='graph')

class dmx.compressor.modeling.model.DmxModelMixin

Bases: object

additional_dmx_aware_mappings: Dict

check_dim_consistency() → bool

A function that checks format dimension consistency and sparseness dimension consistency for all applicable dmx modules in the model

Returns:: True if all dimensions are consistent for all applicable modules.

configure(config: dict | str | None, *rules)

Configure Dmx-specific numerics/sparsity/logics

NOTE: only staticly declared DmxModule(s) are to be transformed

Parameters:

config (Optional[Union[DmxConfig, str]]) – DmxConfig to be used for transformation.
*rules (List[DmxConfigRule]) – variable length of list of configuration rules on top of config.

Returns:

Returns the transformed model

counting_flops(zero: bool = True)

classmethod create_submod_transform_forward(model: Module, submod_name: str, additional_dmx_aware_mappings=None, export=False): Only supported for fx path, submodule forward can be directly called in export path

static deepcopy_args(args)

property dmx_config: “Returns the DmxConfig object for the model

property dmx_module_names: “Returns a list of module names listed in a dmx_config

fold_weights_and_biases(): A function that applies the ops the weights and biases using the corresponding formats.

forward_weight_hypernets()

freeze(config_file='./config.yaml')

A function that stores the state and ops format of the model to a config file

Parameters:: config_file (Optional[str]) – Path to store the config file. Defaults to “./config.yaml”.

get_monitoring_records(submodules_to_monitor: List[str] = [])

get_runtime_records()

static is_same_signature(_model, args, kwargs)

keep_dmx_config()

measure_runtimes(device, submodules_to_measure: List[str] = [])

monitoring(submodules_to_monitor: List[str] = [], save_checkpoint_to: str | None = None)

named_dmx_modules(): “Returns a list of named modules that are dmx configurable

property op_set: Returns a set of unique ops present in the model

static post_process_gm(_model, args, kwargs)

print_model_tree(include_type=False)

A function that prints out the tree structure of a model

Parameters:: include_type (bool) – include the type of modules in the print out if True

static process_inputs_for_export(model, args, kwargs)

thaw(config_file='./config.yaml')

A function that transforms the model in place from a config file.

Parameters:: config_file (Optional[str]) – Path of config file to transform the model from. Defaults to “./config.yaml”.

to_baseline_mode()

to_basic_mode(sbfp_weight_storage=False)

Configures a transformed DmxModel to the BASIC mode on dmx hardware.

Returns:: The configured model.

to_fp8_mode()

Configures a transformed DmxModel to the FP8 mode on dmx hardware.

Returns:: The configured model.

static to_old_forward(_m)

static to_signature_key(_m, _args, _kwargs)

static to_transformed_forward(_m)

transform(config: dict | str | None, *rules)

Configure Dmx-specific numerics/sparsity/logics

NOTE: only staticly declared DmxModule(s) are to be transformed

Parameters:

config (Optional[Union[DmxConfig, str]]) – DmxConfig to be used for transformation.
*rules (List[DmxConfigRule]) – variable length of list of configuration rules on top of config.

Returns:

Returns the transformed model

transformed: bool

class dmx.compressor.modeling.model.DmxPipelineMixin

Bases: object

configure(dmx_config_dict: Dict[str, DmxConfig | None], dmx_transformation_dict: Dict[str, DmxConfigRule | None] | None = None) → None

counting_flops(zero: bool = True) → None

property dmx_config_dict: Dict[str, DmxConfig]

eval()

get_model_by_name(model_name: str) → Module

named_dmx_modules()

transform(dmx_config_dict: Dict[str, DmxConfig | None], dmx_transformation_dict: Dict[str, DmxConfigRule | None] | None = None) → None

class dmx.compressor.modeling.model.DmxSimplePipeline(model_dict: OrderedDict, preproc=Identity(), postproc=Identity(), hf: bool = False, concrete_args_dict: Dict | None = None)

Bases: DmxPipelineMixin

named_dmx_models(): Returns a generator of named DmxModel instances

property op_set

to(torch_device: str | device | None = None)

dmx.compressor.modeling.model.DmxTransformation: alias of DmxConfigRule

class dmx.compressor.modeling.model.Model(body, head=Identity(), tail=Identity(), hf: bool = False, concrete_args: Dict[str, Any] | None = None, **kwargs)

Bases: DmxSimplePipeline

This is a backward-compatible placeholder for legacy models. It is not recommended to use this container; instead, use DmxSimplePipeline directly. TODO: to be deprecated

property op_set: Returns a set of unique ops present in the model