aitune.torch.module.wrapper_module | NVIDIA AITune Documentation

AITune wrapper module.

Module Contents

Classes

Name	Description
`Module`	AITune module wrapper.
`ModuleState`	Possible states of the Module class.

Functions

Name	Description
`get_object_name`	Get the name of an object from its module and class.
`sanitize_model_name`	Sanitize model name to be used as a module name.

Data

API

class aitune.torch.module.wrapper_module.Module(
    module: torch.nn.Module,
    name: str | None = None,
    strategy: aitune.torch.tune_strategy.tune_strategy.TuneStrategy | None = None,
    strategies: aitune.torch.module.wrapper_module.StrategyList | aitune.torch.module.wrapper_module.StrategyMap | None = None,
    dynamic_shapes: aitune.torch.dynamic_shapes.DynamicShapes | None = None
)

Bases: CallableObjectProxy

AITune module wrapper.

This class wraps a torch module and provides tuning functionality. The module can be in 3 different states:

passthrough: the module is not tuned, and will behave identically to the original module.
recording: the module records samples and detects multi-graphs, this is a necessary step before tuning.
tuned: the module is tuned and uses underlying tuned module.

You can go from passthrough and recording with enablePassthrough and enableRecording methods.

You can go from tuned to passthrough/recording only if the force flag is True - then the module will be reset.

This wrapper can be used in place of a torch module, and will behave identically to the original module.

DYNAMIC_SHAPES_KEY

= 'dynamic_shapes'

NAME_KEY

= 'name'

TUNED_MODULE_KEY

= 'tuned_module'

TYPE_KEY

= 'type'

_current_forward_hooks

= module._forward_hooks

_current_forward_pre_hooks

= module._forward_pre_hooks

_self_device

= get_module_device(module) or torch.device('cpu')

_self_name

_self_orig_forward

= module.forward

_self_proxy_forward

= wrapt.decorator(self._forward)(module.forward)

_self_state

= ModuleState.INIT

device

Get the device of the module.

graph_specs

list[GraphSpec]

Multi-graphs of the module.

module

Get the backends of the module.

name

str

Get the name of the module.

state

ModuleState

Get the state of the module.

aitune.torch.module.wrapper_module.Module.__getitem__(
    key: typing.Any
) -> typing.Any

Delegate getitem calls to the wrapped module.

This allows the proxy to handle indexing operations on the wrapped module, which is particularly useful for Sequential modules and other indexable modules.

Parameters:

key

Any

The index or key to use for accessing the wrapped module.

Returns: Any

The result of accessing the wrapped module with the given key.

aitune.torch.module.wrapper_module.Module._activate_wrapper()

aitune.torch.module.wrapper_module.Module._create_graph_cache_dir(
    graph_spec: aitune.torch.module.graph_spec.GraphSpec
) -> pathlib.Path

Create a cache directory for the graph.

aitune.torch.module.wrapper_module.Module._deactivate_wrapper()

aitune.torch.module.wrapper_module.Module._deploy_wrapper(
    device: torch.device | None
)

aitune.torch.module.wrapper_module.Module._forward(
    wrapped,
    instance,
    args,
    kwargs
) -> typing.Any

Calls one of the wrappers depending on the module state.

Before calling a particular wrapper, the forward method is restored so that we avoid infinite recursion. After the implementation handles the call forward is restored back to point to this _forward to intercept subsequent calls.

Parameters:

wrapped

The wrapped module.

instance

The instance of the module.

args

The arguments to pass to the module.

kwargs

The keyword arguments to pass to the module.

Returns: Any

The result of the call.

aitune.torch.module.wrapper_module.Module._get_strategies_for_graph_specs(
    strategy: aitune.torch.tune_strategy.tune_strategy.TuneStrategy | None,
    graph_specs: list[aitune.torch.module.graph_spec.GraphSpec],
    dry_run: bool
) -> list[aitune.torch.tune_strategy.tune_strategy.TuneStrategy]

Returns strategies for given graph specs.

Module level strategies (list/dict) take precedence over provided strategy. The function checks if there is sufficient strategies (list/dict).

aitune.torch.module.wrapper_module.Module._handle_backend_added_hooks()

Handle new hooks if added by a backend.

Before tuning hooks are cleared (empty OrderedDict). If there are new hooks added by a backend, they should be added to the existing ones but before the original hooks so that application layer hooks (like HF post processing hooks) are called after the backend hooks.

aitune.torch.module.wrapper_module.Module._offload(
    backends: collections.OrderedDict[aitune.torch.module.sample_metadata.SampleMetadata, aitune.torch.backend.backend.Backend]
)

Offload the module to the meta device only if no JIT backends are used.

aitune.torch.module.wrapper_module.Module._proxy_forward()

Proxy the forward calls.

We need to re-enable hooks, so that they will be called before and after proxied forward.

aitune.torch.module.wrapper_module.Module._reset()

Resets the module to initial state.

aitune.torch.module.wrapper_module.Module._restore_original_forward()

Restore the original forward and hooks.

We need to disable hooks, otherwise they will be called twice.

aitune.torch.module.wrapper_module.Module._setup_strategies(
    strategy: aitune.torch.tune_strategy.tune_strategy.TuneStrategy | None,
    strategies: aitune.torch.module.wrapper_module.StrategyList | aitune.torch.module.wrapper_module.StrategyMap | None
)

Sets up strategy or strategy_map or strategy_list depending on input args.

aitune.torch.module.wrapper_module.Module.activate()

Activates the module backends.

aitune.torch.module.wrapper_module.Module.deactivate()

Deactivates the module backends.

aitune.torch.module.wrapper_module.Module.enable_passthrough(
    force: bool = False
)

Enables passthrough mode.

Parameters:

force

boolDefaults to False

if True, force the module to be in the passthrough mode.

aitune.torch.module.wrapper_module.Module.enable_recording(
    force: bool = False
)

Enables recording mode.

Parameters:

force

boolDefaults to False

if True, force the module to be in the recording mode.

aitune.torch.module.wrapper_module.Module.from_dict(
    module: torch.nn.Module,
    state_dict: dict,
    device: torch.device | None = None
)

staticmethod

Create a wrapper module from a state dictionary.

aitune.torch.module.wrapper_module.Module.is_state_dict_valid(
    state_dict: dict
)

staticmethod

Check if the state dictionary has a wrapper module.

aitune.torch.module.wrapper_module.Module.to_dict(
    args = (),
    destination = None,
    prefix = '',
    keep_vars = False
)

Convert the wrapper module to a state dictionary.

The method follows torch.state_dict format so that it could be called recursively from torch.nn.Module.state_dict. The name of the function follows aitune convention of from/to dict. Special alias is created to match torch convention i.e. state_dict.

aitune.torch.module.wrapper_module.Module.tune(
    device: torch.device | None = None,
    strategy: aitune.torch.tune_strategy.tune_strategy.TuneStrategy | None = None,
    dry_run: bool = False
)

Tunes the module.

Parameters:

device

torch.device | NoneDefaults to None

the device to use for tuning

strategy

TuneStrategy | NoneDefaults to None

the tuning strategy to use

dry_run

boolDefaults to False

if True, only dry run the tuning

class aitune.torch.module.wrapper_module.ModuleState

Bases: enum.Enum

Possible states of the Module class.

INIT

= 'init'

PASSTHROUGH

= 'passthrough'

RECORDING

= 'recording'

TUNED

= 'tuned'

aitune.torch.module.wrapper_module.get_object_name(
    obj: typing.Any
) -> str

Get the name of an object from its module and class.

aitune.torch.module.wrapper_module.sanitize_model_name(
    model_name: str | None
) -> str | None

Sanitize model name to be used as a module name.

aitune.torch.module.wrapper_module.DEFAULT_STRATEGY = FirstWinsStrategy()

aitune.torch.module.wrapper_module.StrategyList = list[TuneStrategy]

aitune.torch.module.wrapper_module.StrategyMap = dict[SampleMetadata, TuneStrategy]

aitune.torch.module.wrapper_module.logger = getLogger(__name__)