aitune.torch.inspecting.module_inspector | NVIDIA AITune Documentation

Module for inspecting PyTorch models and tracking their execution.

This module is used to inspect PyTorch models and track their execution. It is used to find the modules that are executed and the modules that are not executed. It is also used to wrap the forward methods of the modules to track the execution time and input/output types.

In ModuleInspector, we use vars(obj) to get the members of an object.

Double references to the same module may cause issues. The inspector takes the first reference to the module that is inspected. So, even if the second reference is used for inference, the first will be returned for wrapping.

Object paths used in ModuleInfo are relative to the root module and are in all dot notation, even for dictionaries and lists. (e.g., ‘.list.0.layers, '.dict.key.layers, .encoder.layers.0.self_attn.q_proj.weight).

Environment variables that could be used:

AITUNE_INSPECT_DEBUG: Whether to enable more verbose debug mode for inspecting. Adds visited nodes, and execution order. Default is False.

AITUNE_INSPECT_DEBUG_RAISE: Whether to raise an error when an error occurs during inspecting - ignored by default. Default is False.

Module Contents

Classes

Name	Description
`InspectContext`	Context for inspecting modules.
`ModuleInspector`	Class for inspecting PyTorch modules and tracking their execution.

Functions

Name	Description
`_get_inspect_debug`	-
`_get_inspect_debug_raise`	-

Data

DEFAULT_INSPECT_DEBUG

DEFAULT_INSPECT_DEBUG_RAISE

DEFAULT_MAX_RECURSION_DEPTH

logger

API

class aitune.torch.inspecting.module_inspector.InspectContext(
    depth: int = 0,
    object_path: str = '',
    module_parent: aitune.torch.inspecting.module_info.ModuleInfo | None = None
)

Dataclass

Context for inspecting modules.

depth

int = 0

module_parent

ModuleInfo | None = None

name

str

Get the name.

object_path

str = ''

aitune.torch.inspecting.module_inspector.InspectContext.clone(
    kwargs = {}
) -> aitune.torch.inspecting.module_inspector.InspectContext

Get the inspect context, with optional changes.

aitune.torch.inspecting.module_inspector.InspectContext.create_module_info(
    module: torch.nn.Module | list | dict | typing.Any
) -> aitune.torch.inspecting.module_info.ModuleInfo

Get the ModuleInfo based on current context and provided object type.

aitune.torch.inspecting.module_inspector.InspectContext.next(
    name: str = '',
    parent: aitune.torch.inspecting.module_info.ModuleInfo | None = None
) -> aitune.torch.inspecting.module_inspector.InspectContext

Get the next inspect context - increment depth and add name to object path.

class aitune.torch.inspecting.module_inspector.ModuleInspector(
    min_depth: int = 0,
    max_depth: int = 5
)

Class for inspecting PyTorch modules and tracking their execution.

_execution_parent_module

list[ModuleInfo] = []

_inspected_objects

set = set()

_module_info

dict[Any, ModuleInfo] = {}

_original_forward

dict[Any, Any] = {}

aitune.torch.inspecting.module_inspector.ModuleInspector._debug_error(
    message: str,
    args: typing.Any = (),
    error: Exception
) -> None

Debug an error.

aitune.torch.inspecting.module_inspector.ModuleInspector._debug_inspecting_functions() -> None

Debug inspecting functions.

aitune.torch.inspecting.module_inspector.ModuleInspector._inspect_dict(
    obj: dict,
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> None

Inspect the values of a dictionary.

Parameters:

obj

dict

The dictionary to inspect

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._inspect_list(
    obj: list,
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> None

Inspect the elements of a list.

Parameters:

obj

list

The list to inspect

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._inspect_members(
    obj: typing.Any,
    members: dict[str, typing.Any],
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> None

Inspect the members of an object.

Parameters:

obj

Any

The object to inspect

members

dict[str, Any]

The members of the object to inspect

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._inspect_module(
    module: torch.nn.Module,
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> None

Start inspecting a module and its submodules.

Parameters:

module

torch.nn.Module

The PyTorch module to inspect

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._inspect_object(
    obj: typing.Any,
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> None

Inspect an object and its members for PyTorch modules.

Parameters:

obj

Any

The object to inspect

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._register_module(
    module: torch.nn.Module,
    context: aitune.torch.inspecting.module_inspector.InspectContext
) -> aitune.torch.inspecting.module_info.ModuleInfo | None

Parameters:

module

torch.nn.Module

The module to register

parent

The parent module info if any

context

InspectContext

The context of the inspection

aitune.torch.inspecting.module_inspector.ModuleInspector._should_skip_member(
    obj: typing.Any,
    member_name: str
) -> bool

Check if a member should be skipped.

aitune.torch.inspecting.module_inspector.ModuleInspector._wrap_debug_inspect(
    method: collections.abc.Callable,
    collection: bool = False
) -> collections.abc.Callable

Wrap a method to add debug logging.

aitune.torch.inspecting.module_inspector.ModuleInspector._wrap_forward_methods(
    module
) -> None

Wrap forward methods of a module and its submodules.

Parameters:

module

The module whose forward method to wrap

aitune.torch.inspecting.module_inspector.ModuleInspector.get_modules() -> list[aitune.torch.inspecting.module_info.ModuleInfo]

Get list of top-level executed modules or their first executed children.

Returns: list[ModuleInfo]

List of modules that were executed at the top level or their first executed children

aitune.torch.inspecting.module_inspector.ModuleInspector.inspect(
    obj: typing.Any
) -> None

Inspect an object and its members for PyTorch modules.

Parameters:

obj

Any

The object to inspect

aitune.torch.inspecting.module_inspector.ModuleInspector.nn_module_base_members() -> list[str]

staticmethod

Get the base class members of a PyTorch module.

Returns: list[str]

List of names

aitune.torch.inspecting.module_inspector.ModuleInspector.reset() -> None

Reset the inspector state.

aitune.torch.inspecting.module_inspector._get_inspect_debug() -> bool

aitune.torch.inspecting.module_inspector._get_inspect_debug_raise() -> bool

aitune.torch.inspecting.module_inspector.DEFAULT_INSPECT_DEBUG = False

aitune.torch.inspecting.module_inspector.DEFAULT_INSPECT_DEBUG_RAISE = False

aitune.torch.inspecting.module_inspector.DEFAULT_MAX_RECURSION_DEPTH = 5

aitune.torch.inspecting.module_inspector.logger = getLogger(__name__)